Old Linalg IR - the one which doesn't have the issue

linalg_elided_old.mlir

#map = affine_map<(d0, d1, d2, d3) -> (0, d1, d2, d3)>
#map1 = affine_map<(d0, d1, d2, d3) -> (d0, d1, d2, d3)>
#map2 = affine_map<(d0, d1, d2, d3) -> (d1)>
#map3 = affine_map<(d0, d1, d2, d3) -> (d0, d1, 0, 0)>
#map4 = affine_map<(d0, d1, d2, d3) -> (0, d1, 0, 0)>
#map5 = affine_map<(d0, d1, d2) -> (d0, d1, d2)>
#map6 = affine_map<(d0, d1, d2) -> (d0, d2, d1)>
#map7 = affine_map<(d0, d1) -> (d0, d1)>
#map8 = affine_map<(d0, d1) -> (d1, d0)>
#map9 = affine_map<(d0, d1, d2) -> (0, d1, d2)>
#map10 = affine_map<(d0, d1, d2) -> (d1, d2)>
#map11 = affine_map<(d0, d1, d2) -> (d2)>
#map12 = affine_map<(d0, d1, d2, d3) -> (d0, d2, d1, d3)>
#map13 = affine_map<(d0, d1, d2) -> (d0, d1, 0)>
#map14 = affine_map<(d0, d1, d2) -> (0, d1, 0)>
#map15 = affine_map<(d0, d1) -> (d1)>
module attributes {torch.debug_module_name = "_lambda"} {
  ml_program.global private mutable @global_seed(dense<0> : tensor<i64>) : tensor<i64>
  func.func @forward(%arg0: tensor<1x4x64x64xf16>) -> tensor<1x3x512x512xf16> {
    %cst = arith.constant dense<[1.574710e-02, -2.030940e-02, -4.635620e-02]> : tensor<3xf16>
    %cst_0 = arith.constant dense_resource<__elided__> : tensor<3x128x3x3xf16>
    %cst_1 = arith.constant dense_resource<__elided__> : tensor<128xf16>
    %cst_2 = arith.constant dense_resource<__elided__> : tensor<128xf16>
    %cst_3 = arith.constant dense_resource<__elided__> : tensor<128xf16>
    %cst_4 = arith.constant dense_resource<__elided__> : tensor<128x128x3x3xf16>
    %cst_5 = arith.constant dense_resource<__elided__> : tensor<128xf16>
    %cst_6 = arith.constant dense_resource<__elided__> : tensor<128xf16>
    %cst_7 = arith.constant dense_resource<__elided__> : tensor<128xf16>
    %cst_8 = arith.constant dense_resource<__elided__> : tensor<128x128x3x3xf16>
    %cst_9 = arith.constant dense_resource<__elided__> : tensor<128xf16>
    %cst_10 = arith.constant dense_resource<__elided__> : tensor<128xf16>
    %cst_11 = arith.constant dense_resource<__elided__> : tensor<128xf16>
    %cst_12 = arith.constant dense_resource<__elided__> : tensor<128x128x3x3xf16>
    %cst_13 = arith.constant dense_resource<__elided__> : tensor<128xf16>
    %cst_14 = arith.constant dense_resource<__elided__> : tensor<128xf16>
    %cst_15 = arith.constant dense_resource<__elided__> : tensor<128xf16>
    %cst_16 = arith.constant dense_resource<__elided__> : tensor<128x128x3x3xf16>
    %cst_17 = arith.constant dense_resource<__elided__> : tensor<128xf16>
    %cst_18 = arith.constant dense_resource<__elided__> : tensor<128xf16>
    %cst_19 = arith.constant dense_resource<__elided__> : tensor<128xf16>
    %cst_20 = arith.constant dense_resource<__elided__> : tensor<128x256x1x1xf16>
    %cst_21 = arith.constant dense_resource<__elided__> : tensor<128xf16>
    %cst_22 = arith.constant dense_resource<__elided__> : tensor<128x128x3x3xf16>
    %cst_23 = arith.constant dense_resource<__elided__> : tensor<128xf16>
    %cst_24 = arith.constant dense_resource<__elided__> : tensor<128xf16>
    %cst_25 = arith.constant dense_resource<__elided__> : tensor<128xf16>
    %cst_26 = arith.constant dense_resource<__elided__> : tensor<128x256x3x3xf16>
    %cst_27 = arith.constant dense_resource<__elided__> : tensor<256xf16>
    %cst_28 = arith.constant dense_resource<__elided__> : tensor<256xf16>
    %cst_29 = arith.constant dense_resource<__elided__> : tensor<256xf16>
    %cst_30 = arith.constant dense_resource<__elided__> : tensor<256x256x3x3xf16>
    %cst_31 = arith.constant dense_resource<__elided__> : tensor<256xf16>
    %cst_32 = arith.constant dense_resource<__elided__> : tensor<256x256x3x3xf16>
    %cst_33 = arith.constant dense_resource<__elided__> : tensor<256xf16>
    %cst_34 = arith.constant dense_resource<__elided__> : tensor<256xf16>
    %cst_35 = arith.constant dense_resource<__elided__> : tensor<256xf16>
    %cst_36 = arith.constant dense_resource<__elided__> : tensor<256x256x3x3xf16>
    %cst_37 = arith.constant dense_resource<__elided__> : tensor<256xf16>
    %cst_38 = arith.constant dense_resource<__elided__> : tensor<256xf16>
    %cst_39 = arith.constant dense_resource<__elided__> : tensor<256xf16>
    %cst_40 = arith.constant dense_resource<__elided__> : tensor<256x256x3x3xf16>
    %cst_41 = arith.constant dense_resource<__elided__> : tensor<256xf16>
    %cst_42 = arith.constant dense_resource<__elided__> : tensor<256xf16>
    %cst_43 = arith.constant dense_resource<__elided__> : tensor<256xf16>
    %cst_44 = arith.constant dense_resource<__elided__> : tensor<256x256x3x3xf16>
    %cst_45 = arith.constant dense_resource<__elided__> : tensor<256xf16>
    %cst_46 = arith.constant dense_resource<__elided__> : tensor<256xf16>
    %cst_47 = arith.constant dense_resource<__elided__> : tensor<256xf16>
    %cst_48 = arith.constant dense_resource<__elided__> : tensor<256x512x1x1xf16>
    %cst_49 = arith.constant dense_resource<__elided__> : tensor<256xf16>
    %cst_50 = arith.constant dense_resource<__elided__> : tensor<256x256x3x3xf16>
    %cst_51 = arith.constant dense_resource<__elided__> : tensor<256xf16>
    %cst_52 = arith.constant dense_resource<__elided__> : tensor<256xf16>
    %cst_53 = arith.constant dense_resource<__elided__> : tensor<256xf16>
    %cst_54 = arith.constant dense_resource<__elided__> : tensor<256x512x3x3xf16>
    %cst_55 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_56 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_57 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_58 = arith.constant dense_resource<__elided__> : tensor<512x512x3x3xf16>
    %cst_59 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_60 = arith.constant dense_resource<__elided__> : tensor<512x512x3x3xf16>
    %cst_61 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_62 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_63 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_64 = arith.constant dense_resource<__elided__> : tensor<512x512x3x3xf16>
    %cst_65 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_66 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_67 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_68 = arith.constant dense_resource<__elided__> : tensor<512x512x3x3xf16>
    %cst_69 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_70 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_71 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_72 = arith.constant dense_resource<__elided__> : tensor<512x512x3x3xf16>
    %cst_73 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_74 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_75 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_76 = arith.constant dense_resource<__elided__> : tensor<512x512x3x3xf16>
    %cst_77 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_78 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_79 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_80 = arith.constant dense_resource<__elided__> : tensor<512x512x3x3xf16>
    %cst_81 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_82 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_83 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_84 = arith.constant dense_resource<__elided__> : tensor<512x512x3x3xf16>
    %cst_85 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_86 = arith.constant dense_resource<__elided__> : tensor<512x512x3x3xf16>
    %cst_87 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_88 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_89 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_90 = arith.constant dense_resource<__elided__> : tensor<512x512x3x3xf16>
    %cst_91 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_92 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_93 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_94 = arith.constant dense_resource<__elided__> : tensor<512x512x3x3xf16>
    %cst_95 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_96 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_97 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_98 = arith.constant dense_resource<__elided__> : tensor<512x512x3x3xf16>
    %cst_99 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_100 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_101 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_102 = arith.constant dense_resource<__elided__> : tensor<512x512x3x3xf16>
    %cst_103 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_104 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_105 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_106 = arith.constant dense_resource<__elided__> : tensor<512x512x3x3xf16>
    %cst_107 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_108 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_109 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_110 = arith.constant dense_resource<__elided__> : tensor<512x512x3x3xf16>
    %cst_111 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_112 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_113 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_114 = arith.constant dense_resource<__elided__> : tensor<512x512x3x3xf16>
    %cst_115 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_116 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_117 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_118 = arith.constant dense_resource<__elided__> : tensor<512x512xf16>
    %cst_119 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_120 = arith.constant dense_resource<__elided__> : tensor<512x512xf16>
    %cst_121 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_122 = arith.constant dense_resource<__elided__> : tensor<512x512xf16>
    %cst_123 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_124 = arith.constant dense_resource<__elided__> : tensor<512x512xf16>
    %cst_125 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_126 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_127 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_128 = arith.constant dense_resource<__elided__> : tensor<512x512x3x3xf16>
    %cst_129 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_130 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_131 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_132 = arith.constant dense_resource<__elided__> : tensor<512x512x3x3xf16>
    %cst_133 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_134 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_135 = arith.constant dense_resource<__elided__> : tensor<512xf16>
    %cst_136 = arith.constant dense_resource<__elided__> : tensor<512x4x3x3xf16>
    %cst_137 = arith.constant dense<[3.207400e-02, -8.428960e-02, -2.432860e-01, 1.315920e-01]> : tensor<4xf16>
    %cst_138 = arith.constant dense_resource<__elided__> : tensor<4x4x1x1xf16>
    %cst_139 = arith.constant 0.000000e+00 : f16
    %cst_140 = arith.constant 0.000000e+00 : f32
    %cst_141 = arith.constant 1.000000e+00 : f16
    %cst_142 = arith.constant -6.550400e+04 : f16
    %cst_143 = arith.constant 0.000000e+00 : f64
    %cst_144 = arith.constant 5.4899807850672522 : f64
    %cst_145 = arith.constant 9.9999999999999995E-7 : f64
    %cst_146 = arith.constant 0.044194173824159216 : f64
    %cst_147 = arith.constant 6.553600e+04 : f64
    %cst_148 = arith.constant 6.553600e+04 : f32
    %cst_149 = arith.constant 2.621440e+05 : f64
    %cst_150 = arith.constant 2.621440e+05 : f32
    %cst_151 = arith.constant 0x4130000000000000 : f64
    %cst_152 = arith.constant 0x49800000 : f32
    %cst_153 = arith.constant 5.242880e+05 : f64
    %cst_154 = arith.constant 5.242880e+05 : f32
    %cst_155 = arith.constant 0x4140000000000000 : f64
    %cst_156 = arith.constant 0x4A000000 : f32
    %c0_i64 = arith.constant 0 : i64
    %c0 = arith.constant 0 : index
    %c1 = arith.constant 1 : index
    %cst_157 = arith.constant 2.000000e+00 : f16
    %cst_158 = arith.constant 5.000000e-01 : f16
    %cst_159 = arith.constant 2.550000e+02 : f16
    %c2 = arith.constant 2 : index
    %0 = tensor.empty() : tensor<1x4x64x64xf16>
    %1 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%arg0 : tensor<1x4x64x64xf16>) outs(%0 : tensor<1x4x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_144 : f64 to f16
      %884 = arith.mulf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x4x64x64xf16>
    %2 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_137 : tensor<4xf16>) outs(%0 : tensor<1x4x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x4x64x64xf16>
    %3 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%1, %cst_138 : tensor<1x4x64x64xf16>, tensor<4x4x1x1xf16>) outs(%2 : tensor<1x4x64x64xf16>) -> tensor<1x4x64x64xf16>
    %padded = tensor.pad %3 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x4x64x64xf16> to tensor<1x4x66x66xf16>
    %4 = tensor.empty() : tensor<1x512x64x64xf16>
    %5 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_135 : tensor<512xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x64x64xf16>
    %6 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded, %cst_136 : tensor<1x4x66x66xf16>, tensor<512x4x3x3xf16>) outs(%5 : tensor<1x512x64x64xf16>) -> tensor<1x512x64x64xf16>
    %collapsed = tensor.collapse_shape %6 [[0], [1], [2, 3]] : tensor<1x512x64x64xf16> into tensor<1x512x4096xf16>
    %expanded = tensor.expand_shape %collapsed [[0], [1, 2], [3]] : tensor<1x512x4096xf16> into tensor<1x32x16x4096xf16>
    %7 = tensor.empty() : tensor<1x32x16x4096xf32>
    %8 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded : tensor<1x32x16x4096xf16>) outs(%7 : tensor<1x32x16x4096xf32>) {
    ^bb0(%in: f16, %out: f32):
      %883 = arith.extf %in : f16 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x16x4096xf32>
    %9 = tensor.empty() : tensor<1x32x16x4096xf64>
    %10 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%8 : tensor<1x32x16x4096xf32>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f32, %out: f64):
      %883 = arith.extf %in : f32 to f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %11 = tensor.empty() : tensor<1x32x1x1xf64>
    %12 = linalg.fill ins(%cst_143 : f64) outs(%11 : tensor<1x32x1x1xf64>) -> tensor<1x32x1x1xf64>
    %13 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%10 : tensor<1x32x16x4096xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %14 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%13 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_147 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %15 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%10, %14 : tensor<1x32x16x4096xf64>, tensor<1x32x1x1xf64>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.subf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %16 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%15, %15 : tensor<1x32x16x4096xf64>, tensor<1x32x16x4096xf64>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.mulf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %17 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%16 : tensor<1x32x16x4096xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %18 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%17 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_147 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %19 = tensor.empty() : tensor<1x32x1x1xf32>
    %20 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%18 : tensor<1x32x1x1xf64>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f64, %out: f32):
      %883 = arith.truncf %in : f64 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %21 = linalg.fill ins(%cst_140 : f32) outs(%19 : tensor<1x32x1x1xf32>) -> tensor<1x32x1x1xf32>
    %22 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%8 : tensor<1x32x16x4096xf32>) outs(%21 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.addf %in, %out : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %23 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%22 : tensor<1x32x1x1xf32>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.divf %in, %cst_148 : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %24 = tensor.empty() : tensor<1x32x1x1xf16>
    %25 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%20 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %26 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%25 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %27 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%23 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %28 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%27 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %29 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%26 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_145 : f64 to f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x32x1x1xf16>
    %30 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%29 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.rsqrt %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %31 = tensor.empty() : tensor<1x32x16x4096xf16>
    %32 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded, %28 : tensor<1x32x16x4096xf16>, tensor<1x32x1x1xf16>) outs(%31 : tensor<1x32x16x4096xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x4096xf16>
    %33 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%32, %30 : tensor<1x32x16x4096xf16>, tensor<1x32x1x1xf16>) outs(%31 : tensor<1x32x16x4096xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x4096xf16>
    %collapsed_160 = tensor.collapse_shape %33 [[0], [1, 2], [3]] : tensor<1x32x16x4096xf16> into tensor<1x512x4096xf16>
    %expanded_161 = tensor.expand_shape %collapsed_160 [[0], [1], [2, 3]] : tensor<1x512x4096xf16> into tensor<1x512x64x64xf16>
    %expanded_162 = tensor.expand_shape %cst_134 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %expanded_163 = tensor.expand_shape %cst_133 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %34 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_161, %expanded_163 : tensor<1x512x64x64xf16>, tensor<1x512x1x1xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %35 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%34, %expanded_162 : tensor<1x512x64x64xf16>, tensor<1x512x1x1xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %36 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%35 : tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.negf %in : f16
      %884 = math.exp %883 : f16
      %885 = arith.addf %884, %cst_141 : f16
      %886 = arith.divf %cst_141, %885 : f16
      linalg.yield %886 : f16
    } -> tensor<1x512x64x64xf16>
    %37 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%36, %35 : tensor<1x512x64x64xf16>, tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %padded_164 = tensor.pad %37 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x512x64x64xf16> to tensor<1x512x66x66xf16>
    %38 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_131 : tensor<512xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x64x64xf16>
    %39 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_164, %cst_132 : tensor<1x512x66x66xf16>, tensor<512x512x3x3xf16>) outs(%38 : tensor<1x512x64x64xf16>) -> tensor<1x512x64x64xf16>
    %collapsed_165 = tensor.collapse_shape %39 [[0], [1], [2, 3]] : tensor<1x512x64x64xf16> into tensor<1x512x4096xf16>
    %expanded_166 = tensor.expand_shape %collapsed_165 [[0], [1, 2], [3]] : tensor<1x512x4096xf16> into tensor<1x32x16x4096xf16>
    %40 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_166 : tensor<1x32x16x4096xf16>) outs(%7 : tensor<1x32x16x4096xf32>) {
    ^bb0(%in: f16, %out: f32):
      %883 = arith.extf %in : f16 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x16x4096xf32>
    %41 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%40 : tensor<1x32x16x4096xf32>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f32, %out: f64):
      %883 = arith.extf %in : f32 to f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %42 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%41 : tensor<1x32x16x4096xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %43 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%42 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_147 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %44 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%41, %43 : tensor<1x32x16x4096xf64>, tensor<1x32x1x1xf64>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.subf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %45 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%44, %44 : tensor<1x32x16x4096xf64>, tensor<1x32x16x4096xf64>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.mulf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %46 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%45 : tensor<1x32x16x4096xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %47 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%46 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_147 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %48 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%47 : tensor<1x32x1x1xf64>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f64, %out: f32):
      %883 = arith.truncf %in : f64 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %49 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%40 : tensor<1x32x16x4096xf32>) outs(%21 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.addf %in, %out : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %50 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%49 : tensor<1x32x1x1xf32>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.divf %in, %cst_148 : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %51 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%48 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %52 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%51 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %53 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%50 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %54 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%53 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %55 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%52 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_145 : f64 to f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x32x1x1xf16>
    %56 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%55 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.rsqrt %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %57 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_166, %54 : tensor<1x32x16x4096xf16>, tensor<1x32x1x1xf16>) outs(%31 : tensor<1x32x16x4096xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x4096xf16>
    %58 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%57, %56 : tensor<1x32x16x4096xf16>, tensor<1x32x1x1xf16>) outs(%31 : tensor<1x32x16x4096xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x4096xf16>
    %collapsed_167 = tensor.collapse_shape %58 [[0], [1, 2], [3]] : tensor<1x32x16x4096xf16> into tensor<1x512x4096xf16>
    %expanded_168 = tensor.expand_shape %collapsed_167 [[0], [1], [2, 3]] : tensor<1x512x4096xf16> into tensor<1x512x64x64xf16>
    %expanded_169 = tensor.expand_shape %cst_130 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %expanded_170 = tensor.expand_shape %cst_129 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %59 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_168, %expanded_170 : tensor<1x512x64x64xf16>, tensor<1x512x1x1xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %60 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%59, %expanded_169 : tensor<1x512x64x64xf16>, tensor<1x512x1x1xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %61 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%60 : tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.negf %in : f16
      %884 = math.exp %883 : f16
      %885 = arith.addf %884, %cst_141 : f16
      %886 = arith.divf %cst_141, %885 : f16
      linalg.yield %886 : f16
    } -> tensor<1x512x64x64xf16>
    %62 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%61, %60 : tensor<1x512x64x64xf16>, tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %padded_171 = tensor.pad %62 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x512x64x64xf16> to tensor<1x512x66x66xf16>
    %63 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_127 : tensor<512xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x64x64xf16>
    %64 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_171, %cst_128 : tensor<1x512x66x66xf16>, tensor<512x512x3x3xf16>) outs(%63 : tensor<1x512x64x64xf16>) -> tensor<1x512x64x64xf16>
    %65 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%6, %64 : tensor<1x512x64x64xf16>, tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %66 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%65 : tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x64x64xf16>
    %collapsed_172 = tensor.collapse_shape %66 [[0], [1], [2, 3]] : tensor<1x512x64x64xf16> into tensor<1x512x4096xf16>
    %expanded_173 = tensor.expand_shape %collapsed_172 [[0], [1, 2], [3]] : tensor<1x512x4096xf16> into tensor<1x32x16x4096xf16>
    %67 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_173 : tensor<1x32x16x4096xf16>) outs(%7 : tensor<1x32x16x4096xf32>) {
    ^bb0(%in: f16, %out: f32):
      %883 = arith.extf %in : f16 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x16x4096xf32>
    %68 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%67 : tensor<1x32x16x4096xf32>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f32, %out: f64):
      %883 = arith.extf %in : f32 to f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %69 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%68 : tensor<1x32x16x4096xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %70 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%69 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_147 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %71 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%68, %70 : tensor<1x32x16x4096xf64>, tensor<1x32x1x1xf64>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.subf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %72 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%71, %71 : tensor<1x32x16x4096xf64>, tensor<1x32x16x4096xf64>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.mulf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %73 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%72 : tensor<1x32x16x4096xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %74 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%73 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_147 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %75 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%74 : tensor<1x32x1x1xf64>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f64, %out: f32):
      %883 = arith.truncf %in : f64 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %76 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%67 : tensor<1x32x16x4096xf32>) outs(%21 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.addf %in, %out : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %77 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%76 : tensor<1x32x1x1xf32>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.divf %in, %cst_148 : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %78 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%75 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %79 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%78 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %80 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%77 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %81 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%80 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %82 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%79 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_145 : f64 to f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x32x1x1xf16>
    %83 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%82 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.rsqrt %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %84 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_173, %81 : tensor<1x32x16x4096xf16>, tensor<1x32x1x1xf16>) outs(%31 : tensor<1x32x16x4096xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x4096xf16>
    %85 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%84, %83 : tensor<1x32x16x4096xf16>, tensor<1x32x1x1xf16>) outs(%31 : tensor<1x32x16x4096xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x4096xf16>
    %collapsed_174 = tensor.collapse_shape %85 [[0], [1, 2], [3]] : tensor<1x32x16x4096xf16> into tensor<1x512x4096xf16>
    %expanded_175 = tensor.expand_shape %collapsed_174 [[0], [1], [2, 3]] : tensor<1x512x4096xf16> into tensor<1x512x64x64xf16>
    %expanded_176 = tensor.expand_shape %cst_126 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %expanded_177 = tensor.expand_shape %cst_125 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %86 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_175, %expanded_177 : tensor<1x512x64x64xf16>, tensor<1x512x1x1xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %87 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%86, %expanded_176 : tensor<1x512x64x64xf16>, tensor<1x512x1x1xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %collapsed_178 = tensor.collapse_shape %87 [[0], [1], [2, 3]] : tensor<1x512x64x64xf16> into tensor<1x512x4096xf16>
    %88 = tensor.empty() : tensor<1x4096x512xf16>
    %89 = linalg.generic {indexing_maps = [#map5, #map6], iterator_types = ["parallel", "parallel", "parallel"]} ins(%collapsed_178 : tensor<1x512x4096xf16>) outs(%88 : tensor<1x4096x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x4096x512xf16>
    %90 = tensor.empty() : tensor<512x512xf16>
    %91 = linalg.generic {indexing_maps = [#map7, #map8], iterator_types = ["parallel", "parallel"]} ins(%cst_124 : tensor<512x512xf16>) outs(%90 : tensor<512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<512x512xf16>
    %92 = linalg.generic {indexing_maps = [#map9, #map5], iterator_types = ["parallel", "parallel", "parallel"]} ins(%89 : tensor<1x4096x512xf16>) outs(%88 : tensor<1x4096x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x4096x512xf16>
    %93 = tensor.empty() : tensor<1x512x512xf16>
    %94 = linalg.generic {indexing_maps = [#map10, #map5], iterator_types = ["parallel", "parallel", "parallel"]} ins(%91 : tensor<512x512xf16>) outs(%93 : tensor<1x512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x512xf16>
    %95 = linalg.fill ins(%cst_139 : f16) outs(%88 : tensor<1x4096x512xf16>) -> tensor<1x4096x512xf16>
    %96 = linalg.batch_matmul ins(%92, %94 : tensor<1x4096x512xf16>, tensor<1x512x512xf16>) outs(%95 : tensor<1x4096x512xf16>) -> tensor<1x4096x512xf16>
    %97 = linalg.generic {indexing_maps = [#map9, #map11, #map5], iterator_types = ["parallel", "parallel", "parallel"]} ins(%96, %cst_123 : tensor<1x4096x512xf16>, tensor<512xf16>) outs(%88 : tensor<1x4096x512xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x4096x512xf16>
    %98 = linalg.generic {indexing_maps = [#map7, #map8], iterator_types = ["parallel", "parallel"]} ins(%cst_122 : tensor<512x512xf16>) outs(%90 : tensor<512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<512x512xf16>
    %99 = linalg.generic {indexing_maps = [#map10, #map5], iterator_types = ["parallel", "parallel", "parallel"]} ins(%98 : tensor<512x512xf16>) outs(%93 : tensor<1x512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x512xf16>
    %100 = linalg.batch_matmul ins(%92, %99 : tensor<1x4096x512xf16>, tensor<1x512x512xf16>) outs(%95 : tensor<1x4096x512xf16>) -> tensor<1x4096x512xf16>
    %101 = linalg.generic {indexing_maps = [#map9, #map11, #map5], iterator_types = ["parallel", "parallel", "parallel"]} ins(%100, %cst_121 : tensor<1x4096x512xf16>, tensor<512xf16>) outs(%88 : tensor<1x4096x512xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x4096x512xf16>
    %102 = linalg.generic {indexing_maps = [#map7, #map8], iterator_types = ["parallel", "parallel"]} ins(%cst_120 : tensor<512x512xf16>) outs(%90 : tensor<512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<512x512xf16>
    %103 = linalg.generic {indexing_maps = [#map10, #map5], iterator_types = ["parallel", "parallel", "parallel"]} ins(%102 : tensor<512x512xf16>) outs(%93 : tensor<1x512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x512xf16>
    %104 = linalg.batch_matmul ins(%92, %103 : tensor<1x4096x512xf16>, tensor<1x512x512xf16>) outs(%95 : tensor<1x4096x512xf16>) -> tensor<1x4096x512xf16>
    %105 = linalg.generic {indexing_maps = [#map9, #map11, #map5], iterator_types = ["parallel", "parallel", "parallel"]} ins(%104, %cst_119 : tensor<1x4096x512xf16>, tensor<512xf16>) outs(%88 : tensor<1x4096x512xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x4096x512xf16>
    %expanded_179 = tensor.expand_shape %97 [[0], [1], [2, 3]] : tensor<1x4096x512xf16> into tensor<1x4096x1x512xf16>
    %106 = tensor.empty() : tensor<1x1x4096x512xf16>
    %107 = linalg.generic {indexing_maps = [#map1, #map12], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_179 : tensor<1x4096x1x512xf16>) outs(%106 : tensor<1x1x4096x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x1x4096x512xf16>
    %collapsed_180 = tensor.collapse_shape %107 [[0], [1, 2], [3]] : tensor<1x1x4096x512xf16> into tensor<1x4096x512xf16>
    %expanded_181 = tensor.expand_shape %101 [[0], [1], [2, 3]] : tensor<1x4096x512xf16> into tensor<1x4096x1x512xf16>
    %108 = linalg.generic {indexing_maps = [#map1, #map12], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_181 : tensor<1x4096x1x512xf16>) outs(%106 : tensor<1x1x4096x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x1x4096x512xf16>
    %collapsed_182 = tensor.collapse_shape %108 [[0], [1, 2], [3]] : tensor<1x1x4096x512xf16> into tensor<1x4096x512xf16>
    %expanded_183 = tensor.expand_shape %105 [[0], [1], [2, 3]] : tensor<1x4096x512xf16> into tensor<1x4096x1x512xf16>
    %109 = linalg.generic {indexing_maps = [#map1, #map12], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_183 : tensor<1x4096x1x512xf16>) outs(%106 : tensor<1x1x4096x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x1x4096x512xf16>
    %collapsed_184 = tensor.collapse_shape %109 [[0], [1, 2], [3]] : tensor<1x1x4096x512xf16> into tensor<1x4096x512xf16>
    %110 = tensor.empty() : tensor<1x4096x4096xf16>
    %111 = linalg.generic {indexing_maps = [#map9, #map5], iterator_types = ["parallel", "parallel", "parallel"]} ins(%110 : tensor<1x4096x4096xf16>) outs(%110 : tensor<1x4096x4096xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %cst_139 : f16
    } -> tensor<1x4096x4096xf16>
    %112 = tensor.empty() : tensor<1x512x4096xf16>
    %113 = linalg.generic {indexing_maps = [#map5, #map6], iterator_types = ["parallel", "parallel", "parallel"]} ins(%collapsed_182 : tensor<1x4096x512xf16>) outs(%112 : tensor<1x512x4096xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x4096xf16>
    %114 = linalg.fill ins(%cst_139 : f16) outs(%110 : tensor<1x4096x4096xf16>) -> tensor<1x4096x4096xf16>
    %115 = linalg.batch_matmul ins(%collapsed_180, %113 : tensor<1x4096x512xf16>, tensor<1x512x4096xf16>) outs(%114 : tensor<1x4096x4096xf16>) -> tensor<1x4096x4096xf16>
    %116 = linalg.generic {indexing_maps = [#map9, #map5], iterator_types = ["parallel", "parallel", "parallel"]} ins(%115 : tensor<1x4096x4096xf16>) outs(%110 : tensor<1x4096x4096xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_146 : f64 to f16
      %884 = arith.mulf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x4096x4096xf16>
    %117 = linalg.generic {indexing_maps = [#map9, #map9, #map5], iterator_types = ["parallel", "parallel", "parallel"]} ins(%116, %111 : tensor<1x4096x4096xf16>, tensor<1x4096x4096xf16>) outs(%110 : tensor<1x4096x4096xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in_381, %cst_139 : f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x4096x4096xf16>
    %118 = tensor.empty() : tensor<1x4096x1xi64>
    %119 = linalg.fill ins(%c0_i64 : i64) outs(%118 : tensor<1x4096x1xi64>) -> tensor<1x4096x1xi64>
    %120 = tensor.empty() : tensor<1x4096x1xf16>
    %121 = linalg.fill ins(%cst_142 : f16) outs(%120 : tensor<1x4096x1xf16>) -> tensor<1x4096x1xf16>
    %122:2 = linalg.generic {indexing_maps = [#map5, #map13, #map13], iterator_types = ["parallel", "parallel", "reduction"]} ins(%117 : tensor<1x4096x4096xf16>) outs(%121, %119 : tensor<1x4096x1xf16>, tensor<1x4096x1xi64>) {
    ^bb0(%in: f16, %out: f16, %out_381: i64):
      %883 = linalg.index 2 : index
      %884 = arith.index_cast %883 : index to i64
      %885 = arith.maxf %in, %out : f16
      %886 = arith.cmpf ogt, %in, %out : f16
      %887 = arith.select %886, %884, %out_381 : i64
      linalg.yield %885, %887 : f16, i64
    } -> (tensor<1x4096x1xf16>, tensor<1x4096x1xi64>)
    %123 = linalg.generic {indexing_maps = [#map9, #map14, #map5], iterator_types = ["parallel", "parallel", "parallel"]} ins(%117, %122#0 : tensor<1x4096x4096xf16>, tensor<1x4096x1xf16>) outs(%110 : tensor<1x4096x4096xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x4096x4096xf16>
    %124 = linalg.generic {indexing_maps = [#map9, #map5], iterator_types = ["parallel", "parallel", "parallel"]} ins(%123 : tensor<1x4096x4096xf16>) outs(%110 : tensor<1x4096x4096xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.exp %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x4096x4096xf16>
    %125 = linalg.fill ins(%cst_139 : f16) outs(%120 : tensor<1x4096x1xf16>) -> tensor<1x4096x1xf16>
    %126 = linalg.generic {indexing_maps = [#map5, #map13], iterator_types = ["parallel", "parallel", "reduction"]} ins(%124 : tensor<1x4096x4096xf16>) outs(%125 : tensor<1x4096x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.addf %in, %out : f16
      linalg.yield %883 : f16
    } -> tensor<1x4096x1xf16>
    %127 = linalg.generic {indexing_maps = [#map9, #map14, #map5], iterator_types = ["parallel", "parallel", "parallel"]} ins(%124, %126 : tensor<1x4096x4096xf16>, tensor<1x4096x1xf16>) outs(%110 : tensor<1x4096x4096xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.divf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x4096x4096xf16>
    %128 = linalg.batch_matmul ins(%127, %collapsed_184 : tensor<1x4096x4096xf16>, tensor<1x4096x512xf16>) outs(%95 : tensor<1x4096x512xf16>) -> tensor<1x4096x512xf16>
    %expanded_185 = tensor.expand_shape %128 [[0], [1, 2], [3]] : tensor<1x4096x512xf16> into tensor<1x1x4096x512xf16>
    %129 = tensor.empty() : tensor<1x4096x1x512xf16>
    %130 = linalg.generic {indexing_maps = [#map1, #map12], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_185 : tensor<1x1x4096x512xf16>) outs(%129 : tensor<1x4096x1x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x4096x1x512xf16>
    %131 = linalg.generic {indexing_maps = [#map7, #map8], iterator_types = ["parallel", "parallel"]} ins(%cst_118 : tensor<512x512xf16>) outs(%90 : tensor<512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<512x512xf16>
    %collapsed_186 = tensor.collapse_shape %130 [[0, 1], [2, 3]] : tensor<1x4096x1x512xf16> into tensor<4096x512xf16>
    %132 = tensor.empty() : tensor<4096x512xf16>
    %133 = linalg.fill ins(%cst_139 : f16) outs(%132 : tensor<4096x512xf16>) -> tensor<4096x512xf16>
    %134 = linalg.matmul ins(%collapsed_186, %131 : tensor<4096x512xf16>, tensor<512x512xf16>) outs(%133 : tensor<4096x512xf16>) -> tensor<4096x512xf16>
    %135 = linalg.generic {indexing_maps = [#map15, #map7, #map7], iterator_types = ["parallel", "parallel"]} ins(%cst_117, %134 : tensor<512xf16>, tensor<4096x512xf16>) outs(%132 : tensor<4096x512xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<4096x512xf16>
    %expanded_187 = tensor.expand_shape %135 [[0, 1], [2]] : tensor<4096x512xf16> into tensor<1x4096x512xf16>
    %136 = linalg.generic {indexing_maps = [#map5, #map6], iterator_types = ["parallel", "parallel", "parallel"]} ins(%expanded_187 : tensor<1x4096x512xf16>) outs(%112 : tensor<1x512x4096xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x4096xf16>
    %expanded_188 = tensor.expand_shape %136 [[0], [1], [2, 3]] : tensor<1x512x4096xf16> into tensor<1x512x64x64xf16>
    %137 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_188, %66 : tensor<1x512x64x64xf16>, tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %138 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%137 : tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x64x64xf16>
    %139 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%138 : tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x64x64xf16>
    %collapsed_189 = tensor.collapse_shape %139 [[0], [1], [2, 3]] : tensor<1x512x64x64xf16> into tensor<1x512x4096xf16>
    %expanded_190 = tensor.expand_shape %collapsed_189 [[0], [1, 2], [3]] : tensor<1x512x4096xf16> into tensor<1x32x16x4096xf16>
    %140 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_190 : tensor<1x32x16x4096xf16>) outs(%7 : tensor<1x32x16x4096xf32>) {
    ^bb0(%in: f16, %out: f32):
      %883 = arith.extf %in : f16 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x16x4096xf32>
    %141 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%140 : tensor<1x32x16x4096xf32>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f32, %out: f64):
      %883 = arith.extf %in : f32 to f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %142 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%141 : tensor<1x32x16x4096xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %143 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%142 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_147 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %144 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%141, %143 : tensor<1x32x16x4096xf64>, tensor<1x32x1x1xf64>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.subf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %145 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%144, %144 : tensor<1x32x16x4096xf64>, tensor<1x32x16x4096xf64>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.mulf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %146 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%145 : tensor<1x32x16x4096xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %147 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%146 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_147 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %148 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%147 : tensor<1x32x1x1xf64>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f64, %out: f32):
      %883 = arith.truncf %in : f64 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %149 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%140 : tensor<1x32x16x4096xf32>) outs(%21 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.addf %in, %out : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %150 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%149 : tensor<1x32x1x1xf32>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.divf %in, %cst_148 : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %151 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%148 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %152 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%151 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %153 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%150 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %154 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%153 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %155 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%152 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_145 : f64 to f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x32x1x1xf16>
    %156 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%155 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.rsqrt %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %157 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_190, %154 : tensor<1x32x16x4096xf16>, tensor<1x32x1x1xf16>) outs(%31 : tensor<1x32x16x4096xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x4096xf16>
    %158 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%157, %156 : tensor<1x32x16x4096xf16>, tensor<1x32x1x1xf16>) outs(%31 : tensor<1x32x16x4096xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x4096xf16>
    %collapsed_191 = tensor.collapse_shape %158 [[0], [1, 2], [3]] : tensor<1x32x16x4096xf16> into tensor<1x512x4096xf16>
    %expanded_192 = tensor.expand_shape %collapsed_191 [[0], [1], [2, 3]] : tensor<1x512x4096xf16> into tensor<1x512x64x64xf16>
    %expanded_193 = tensor.expand_shape %cst_116 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %expanded_194 = tensor.expand_shape %cst_115 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %159 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_192, %expanded_194 : tensor<1x512x64x64xf16>, tensor<1x512x1x1xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %160 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%159, %expanded_193 : tensor<1x512x64x64xf16>, tensor<1x512x1x1xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %161 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%160 : tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.negf %in : f16
      %884 = math.exp %883 : f16
      %885 = arith.addf %884, %cst_141 : f16
      %886 = arith.divf %cst_141, %885 : f16
      linalg.yield %886 : f16
    } -> tensor<1x512x64x64xf16>
    %162 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%161, %160 : tensor<1x512x64x64xf16>, tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %padded_195 = tensor.pad %162 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x512x64x64xf16> to tensor<1x512x66x66xf16>
    %163 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_113 : tensor<512xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x64x64xf16>
    %164 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_195, %cst_114 : tensor<1x512x66x66xf16>, tensor<512x512x3x3xf16>) outs(%163 : tensor<1x512x64x64xf16>) -> tensor<1x512x64x64xf16>
    %collapsed_196 = tensor.collapse_shape %164 [[0], [1], [2, 3]] : tensor<1x512x64x64xf16> into tensor<1x512x4096xf16>
    %expanded_197 = tensor.expand_shape %collapsed_196 [[0], [1, 2], [3]] : tensor<1x512x4096xf16> into tensor<1x32x16x4096xf16>
    %165 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_197 : tensor<1x32x16x4096xf16>) outs(%7 : tensor<1x32x16x4096xf32>) {
    ^bb0(%in: f16, %out: f32):
      %883 = arith.extf %in : f16 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x16x4096xf32>
    %166 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%165 : tensor<1x32x16x4096xf32>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f32, %out: f64):
      %883 = arith.extf %in : f32 to f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %167 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%166 : tensor<1x32x16x4096xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %168 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%167 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_147 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %169 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%166, %168 : tensor<1x32x16x4096xf64>, tensor<1x32x1x1xf64>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.subf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %170 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%169, %169 : tensor<1x32x16x4096xf64>, tensor<1x32x16x4096xf64>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.mulf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %171 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%170 : tensor<1x32x16x4096xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %172 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%171 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_147 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %173 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%172 : tensor<1x32x1x1xf64>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f64, %out: f32):
      %883 = arith.truncf %in : f64 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %174 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%165 : tensor<1x32x16x4096xf32>) outs(%21 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.addf %in, %out : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %175 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%174 : tensor<1x32x1x1xf32>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.divf %in, %cst_148 : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %176 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%173 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %177 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%176 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %178 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%175 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %179 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%178 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %180 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%177 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_145 : f64 to f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x32x1x1xf16>
    %181 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%180 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.rsqrt %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %182 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_197, %179 : tensor<1x32x16x4096xf16>, tensor<1x32x1x1xf16>) outs(%31 : tensor<1x32x16x4096xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x4096xf16>
    %183 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%182, %181 : tensor<1x32x16x4096xf16>, tensor<1x32x1x1xf16>) outs(%31 : tensor<1x32x16x4096xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x4096xf16>
    %collapsed_198 = tensor.collapse_shape %183 [[0], [1, 2], [3]] : tensor<1x32x16x4096xf16> into tensor<1x512x4096xf16>
    %expanded_199 = tensor.expand_shape %collapsed_198 [[0], [1], [2, 3]] : tensor<1x512x4096xf16> into tensor<1x512x64x64xf16>
    %expanded_200 = tensor.expand_shape %cst_112 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %expanded_201 = tensor.expand_shape %cst_111 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %184 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_199, %expanded_201 : tensor<1x512x64x64xf16>, tensor<1x512x1x1xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %185 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%184, %expanded_200 : tensor<1x512x64x64xf16>, tensor<1x512x1x1xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %186 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%185 : tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.negf %in : f16
      %884 = math.exp %883 : f16
      %885 = arith.addf %884, %cst_141 : f16
      %886 = arith.divf %cst_141, %885 : f16
      linalg.yield %886 : f16
    } -> tensor<1x512x64x64xf16>
    %187 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%186, %185 : tensor<1x512x64x64xf16>, tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %padded_202 = tensor.pad %187 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x512x64x64xf16> to tensor<1x512x66x66xf16>
    %188 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_109 : tensor<512xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x64x64xf16>
    %189 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_202, %cst_110 : tensor<1x512x66x66xf16>, tensor<512x512x3x3xf16>) outs(%188 : tensor<1x512x64x64xf16>) -> tensor<1x512x64x64xf16>
    %190 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%138, %189 : tensor<1x512x64x64xf16>, tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %191 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%190 : tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x64x64xf16>
    %192 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%191 : tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x64x64xf16>
    %collapsed_203 = tensor.collapse_shape %192 [[0], [1], [2, 3]] : tensor<1x512x64x64xf16> into tensor<1x512x4096xf16>
    %expanded_204 = tensor.expand_shape %collapsed_203 [[0], [1, 2], [3]] : tensor<1x512x4096xf16> into tensor<1x32x16x4096xf16>
    %193 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_204 : tensor<1x32x16x4096xf16>) outs(%7 : tensor<1x32x16x4096xf32>) {
    ^bb0(%in: f16, %out: f32):
      %883 = arith.extf %in : f16 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x16x4096xf32>
    %194 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%193 : tensor<1x32x16x4096xf32>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f32, %out: f64):
      %883 = arith.extf %in : f32 to f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %195 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%194 : tensor<1x32x16x4096xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %196 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%195 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_147 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %197 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%194, %196 : tensor<1x32x16x4096xf64>, tensor<1x32x1x1xf64>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.subf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %198 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%197, %197 : tensor<1x32x16x4096xf64>, tensor<1x32x16x4096xf64>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.mulf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %199 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%198 : tensor<1x32x16x4096xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %200 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%199 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_147 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %201 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%200 : tensor<1x32x1x1xf64>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f64, %out: f32):
      %883 = arith.truncf %in : f64 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %202 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%193 : tensor<1x32x16x4096xf32>) outs(%21 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.addf %in, %out : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %203 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%202 : tensor<1x32x1x1xf32>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.divf %in, %cst_148 : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %204 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%201 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %205 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%204 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %206 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%203 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %207 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%206 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %208 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%205 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_145 : f64 to f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x32x1x1xf16>
    %209 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%208 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.rsqrt %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %210 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_204, %207 : tensor<1x32x16x4096xf16>, tensor<1x32x1x1xf16>) outs(%31 : tensor<1x32x16x4096xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x4096xf16>
    %211 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%210, %209 : tensor<1x32x16x4096xf16>, tensor<1x32x1x1xf16>) outs(%31 : tensor<1x32x16x4096xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x4096xf16>
    %collapsed_205 = tensor.collapse_shape %211 [[0], [1, 2], [3]] : tensor<1x32x16x4096xf16> into tensor<1x512x4096xf16>
    %expanded_206 = tensor.expand_shape %collapsed_205 [[0], [1], [2, 3]] : tensor<1x512x4096xf16> into tensor<1x512x64x64xf16>
    %expanded_207 = tensor.expand_shape %cst_108 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %expanded_208 = tensor.expand_shape %cst_107 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %212 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_206, %expanded_208 : tensor<1x512x64x64xf16>, tensor<1x512x1x1xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %213 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%212, %expanded_207 : tensor<1x512x64x64xf16>, tensor<1x512x1x1xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %214 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%213 : tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.negf %in : f16
      %884 = math.exp %883 : f16
      %885 = arith.addf %884, %cst_141 : f16
      %886 = arith.divf %cst_141, %885 : f16
      linalg.yield %886 : f16
    } -> tensor<1x512x64x64xf16>
    %215 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%214, %213 : tensor<1x512x64x64xf16>, tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %padded_209 = tensor.pad %215 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x512x64x64xf16> to tensor<1x512x66x66xf16>
    %216 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_105 : tensor<512xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x64x64xf16>
    %217 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_209, %cst_106 : tensor<1x512x66x66xf16>, tensor<512x512x3x3xf16>) outs(%216 : tensor<1x512x64x64xf16>) -> tensor<1x512x64x64xf16>
    %collapsed_210 = tensor.collapse_shape %217 [[0], [1], [2, 3]] : tensor<1x512x64x64xf16> into tensor<1x512x4096xf16>
    %expanded_211 = tensor.expand_shape %collapsed_210 [[0], [1, 2], [3]] : tensor<1x512x4096xf16> into tensor<1x32x16x4096xf16>
    %218 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_211 : tensor<1x32x16x4096xf16>) outs(%7 : tensor<1x32x16x4096xf32>) {
    ^bb0(%in: f16, %out: f32):
      %883 = arith.extf %in : f16 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x16x4096xf32>
    %219 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%218 : tensor<1x32x16x4096xf32>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f32, %out: f64):
      %883 = arith.extf %in : f32 to f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %220 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%219 : tensor<1x32x16x4096xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %221 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%220 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_147 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %222 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%219, %221 : tensor<1x32x16x4096xf64>, tensor<1x32x1x1xf64>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.subf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %223 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%222, %222 : tensor<1x32x16x4096xf64>, tensor<1x32x16x4096xf64>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.mulf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %224 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%223 : tensor<1x32x16x4096xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %225 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%224 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_147 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %226 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%225 : tensor<1x32x1x1xf64>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f64, %out: f32):
      %883 = arith.truncf %in : f64 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %227 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%218 : tensor<1x32x16x4096xf32>) outs(%21 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.addf %in, %out : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %228 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%227 : tensor<1x32x1x1xf32>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.divf %in, %cst_148 : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %229 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%226 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %230 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%229 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %231 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%228 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %232 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%231 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %233 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%230 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_145 : f64 to f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x32x1x1xf16>
    %234 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%233 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.rsqrt %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %235 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_211, %232 : tensor<1x32x16x4096xf16>, tensor<1x32x1x1xf16>) outs(%31 : tensor<1x32x16x4096xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x4096xf16>
    %236 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%235, %234 : tensor<1x32x16x4096xf16>, tensor<1x32x1x1xf16>) outs(%31 : tensor<1x32x16x4096xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x4096xf16>
    %collapsed_212 = tensor.collapse_shape %236 [[0], [1, 2], [3]] : tensor<1x32x16x4096xf16> into tensor<1x512x4096xf16>
    %expanded_213 = tensor.expand_shape %collapsed_212 [[0], [1], [2, 3]] : tensor<1x512x4096xf16> into tensor<1x512x64x64xf16>
    %expanded_214 = tensor.expand_shape %cst_104 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %expanded_215 = tensor.expand_shape %cst_103 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %237 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_213, %expanded_215 : tensor<1x512x64x64xf16>, tensor<1x512x1x1xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %238 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%237, %expanded_214 : tensor<1x512x64x64xf16>, tensor<1x512x1x1xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %239 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%238 : tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.negf %in : f16
      %884 = math.exp %883 : f16
      %885 = arith.addf %884, %cst_141 : f16
      %886 = arith.divf %cst_141, %885 : f16
      linalg.yield %886 : f16
    } -> tensor<1x512x64x64xf16>
    %240 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%239, %238 : tensor<1x512x64x64xf16>, tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %padded_216 = tensor.pad %240 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x512x64x64xf16> to tensor<1x512x66x66xf16>
    %241 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_101 : tensor<512xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x64x64xf16>
    %242 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_216, %cst_102 : tensor<1x512x66x66xf16>, tensor<512x512x3x3xf16>) outs(%241 : tensor<1x512x64x64xf16>) -> tensor<1x512x64x64xf16>
    %243 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%191, %242 : tensor<1x512x64x64xf16>, tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %244 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%243 : tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x64x64xf16>
    %245 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%244 : tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x64x64xf16>
    %collapsed_217 = tensor.collapse_shape %245 [[0], [1], [2, 3]] : tensor<1x512x64x64xf16> into tensor<1x512x4096xf16>
    %expanded_218 = tensor.expand_shape %collapsed_217 [[0], [1, 2], [3]] : tensor<1x512x4096xf16> into tensor<1x32x16x4096xf16>
    %246 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_218 : tensor<1x32x16x4096xf16>) outs(%7 : tensor<1x32x16x4096xf32>) {
    ^bb0(%in: f16, %out: f32):
      %883 = arith.extf %in : f16 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x16x4096xf32>
    %247 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%246 : tensor<1x32x16x4096xf32>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f32, %out: f64):
      %883 = arith.extf %in : f32 to f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %248 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%247 : tensor<1x32x16x4096xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %249 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%248 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_147 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %250 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%247, %249 : tensor<1x32x16x4096xf64>, tensor<1x32x1x1xf64>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.subf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %251 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%250, %250 : tensor<1x32x16x4096xf64>, tensor<1x32x16x4096xf64>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.mulf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %252 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%251 : tensor<1x32x16x4096xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %253 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%252 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_147 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %254 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%253 : tensor<1x32x1x1xf64>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f64, %out: f32):
      %883 = arith.truncf %in : f64 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %255 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%246 : tensor<1x32x16x4096xf32>) outs(%21 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.addf %in, %out : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %256 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%255 : tensor<1x32x1x1xf32>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.divf %in, %cst_148 : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %257 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%254 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %258 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%257 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %259 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%256 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %260 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%259 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %261 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%258 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_145 : f64 to f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x32x1x1xf16>
    %262 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%261 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.rsqrt %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %263 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_218, %260 : tensor<1x32x16x4096xf16>, tensor<1x32x1x1xf16>) outs(%31 : tensor<1x32x16x4096xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x4096xf16>
    %264 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%263, %262 : tensor<1x32x16x4096xf16>, tensor<1x32x1x1xf16>) outs(%31 : tensor<1x32x16x4096xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x4096xf16>
    %collapsed_219 = tensor.collapse_shape %264 [[0], [1, 2], [3]] : tensor<1x32x16x4096xf16> into tensor<1x512x4096xf16>
    %expanded_220 = tensor.expand_shape %collapsed_219 [[0], [1], [2, 3]] : tensor<1x512x4096xf16> into tensor<1x512x64x64xf16>
    %expanded_221 = tensor.expand_shape %cst_100 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %expanded_222 = tensor.expand_shape %cst_99 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %265 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_220, %expanded_222 : tensor<1x512x64x64xf16>, tensor<1x512x1x1xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %266 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%265, %expanded_221 : tensor<1x512x64x64xf16>, tensor<1x512x1x1xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %267 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%266 : tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.negf %in : f16
      %884 = math.exp %883 : f16
      %885 = arith.addf %884, %cst_141 : f16
      %886 = arith.divf %cst_141, %885 : f16
      linalg.yield %886 : f16
    } -> tensor<1x512x64x64xf16>
    %268 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%267, %266 : tensor<1x512x64x64xf16>, tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %padded_223 = tensor.pad %268 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x512x64x64xf16> to tensor<1x512x66x66xf16>
    %269 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_97 : tensor<512xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x64x64xf16>
    %270 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_223, %cst_98 : tensor<1x512x66x66xf16>, tensor<512x512x3x3xf16>) outs(%269 : tensor<1x512x64x64xf16>) -> tensor<1x512x64x64xf16>
    %collapsed_224 = tensor.collapse_shape %270 [[0], [1], [2, 3]] : tensor<1x512x64x64xf16> into tensor<1x512x4096xf16>
    %expanded_225 = tensor.expand_shape %collapsed_224 [[0], [1, 2], [3]] : tensor<1x512x4096xf16> into tensor<1x32x16x4096xf16>
    %271 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_225 : tensor<1x32x16x4096xf16>) outs(%7 : tensor<1x32x16x4096xf32>) {
    ^bb0(%in: f16, %out: f32):
      %883 = arith.extf %in : f16 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x16x4096xf32>
    %272 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%271 : tensor<1x32x16x4096xf32>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f32, %out: f64):
      %883 = arith.extf %in : f32 to f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %273 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%272 : tensor<1x32x16x4096xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %274 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%273 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_147 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %275 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%272, %274 : tensor<1x32x16x4096xf64>, tensor<1x32x1x1xf64>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.subf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %276 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%275, %275 : tensor<1x32x16x4096xf64>, tensor<1x32x16x4096xf64>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.mulf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %277 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%276 : tensor<1x32x16x4096xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %278 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%277 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_147 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %279 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%278 : tensor<1x32x1x1xf64>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f64, %out: f32):
      %883 = arith.truncf %in : f64 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %280 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%271 : tensor<1x32x16x4096xf32>) outs(%21 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.addf %in, %out : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %281 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%280 : tensor<1x32x1x1xf32>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.divf %in, %cst_148 : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %282 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%279 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %283 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%282 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %284 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%281 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %285 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%284 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %286 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%283 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_145 : f64 to f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x32x1x1xf16>
    %287 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%286 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.rsqrt %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %288 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_225, %285 : tensor<1x32x16x4096xf16>, tensor<1x32x1x1xf16>) outs(%31 : tensor<1x32x16x4096xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x4096xf16>
    %289 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%288, %287 : tensor<1x32x16x4096xf16>, tensor<1x32x1x1xf16>) outs(%31 : tensor<1x32x16x4096xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x4096xf16>
    %collapsed_226 = tensor.collapse_shape %289 [[0], [1, 2], [3]] : tensor<1x32x16x4096xf16> into tensor<1x512x4096xf16>
    %expanded_227 = tensor.expand_shape %collapsed_226 [[0], [1], [2, 3]] : tensor<1x512x4096xf16> into tensor<1x512x64x64xf16>
    %expanded_228 = tensor.expand_shape %cst_96 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %expanded_229 = tensor.expand_shape %cst_95 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %290 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_227, %expanded_229 : tensor<1x512x64x64xf16>, tensor<1x512x1x1xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %291 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%290, %expanded_228 : tensor<1x512x64x64xf16>, tensor<1x512x1x1xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %292 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%291 : tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.negf %in : f16
      %884 = math.exp %883 : f16
      %885 = arith.addf %884, %cst_141 : f16
      %886 = arith.divf %cst_141, %885 : f16
      linalg.yield %886 : f16
    } -> tensor<1x512x64x64xf16>
    %293 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%292, %291 : tensor<1x512x64x64xf16>, tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %padded_230 = tensor.pad %293 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x512x64x64xf16> to tensor<1x512x66x66xf16>
    %294 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_93 : tensor<512xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x64x64xf16>
    %295 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_230, %cst_94 : tensor<1x512x66x66xf16>, tensor<512x512x3x3xf16>) outs(%294 : tensor<1x512x64x64xf16>) -> tensor<1x512x64x64xf16>
    %296 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%244, %295 : tensor<1x512x64x64xf16>, tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %297 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%296 : tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x64x64xf16>
    %298 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%297 : tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x64x64xf16>
    %collapsed_231 = tensor.collapse_shape %298 [[0], [1], [2, 3]] : tensor<1x512x64x64xf16> into tensor<1x512x4096xf16>
    %expanded_232 = tensor.expand_shape %collapsed_231 [[0], [1, 2], [3]] : tensor<1x512x4096xf16> into tensor<1x32x16x4096xf16>
    %299 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_232 : tensor<1x32x16x4096xf16>) outs(%7 : tensor<1x32x16x4096xf32>) {
    ^bb0(%in: f16, %out: f32):
      %883 = arith.extf %in : f16 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x16x4096xf32>
    %300 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%299 : tensor<1x32x16x4096xf32>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f32, %out: f64):
      %883 = arith.extf %in : f32 to f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %301 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%300 : tensor<1x32x16x4096xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %302 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%301 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_147 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %303 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%300, %302 : tensor<1x32x16x4096xf64>, tensor<1x32x1x1xf64>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.subf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %304 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%303, %303 : tensor<1x32x16x4096xf64>, tensor<1x32x16x4096xf64>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.mulf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %305 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%304 : tensor<1x32x16x4096xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %306 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%305 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_147 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %307 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%306 : tensor<1x32x1x1xf64>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f64, %out: f32):
      %883 = arith.truncf %in : f64 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %308 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%299 : tensor<1x32x16x4096xf32>) outs(%21 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.addf %in, %out : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %309 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%308 : tensor<1x32x1x1xf32>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.divf %in, %cst_148 : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %310 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%307 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %311 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%310 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %312 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%309 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %313 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%312 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %314 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%311 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_145 : f64 to f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x32x1x1xf16>
    %315 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%314 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.rsqrt %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %316 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_232, %313 : tensor<1x32x16x4096xf16>, tensor<1x32x1x1xf16>) outs(%31 : tensor<1x32x16x4096xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x4096xf16>
    %317 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%316, %315 : tensor<1x32x16x4096xf16>, tensor<1x32x1x1xf16>) outs(%31 : tensor<1x32x16x4096xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x4096xf16>
    %collapsed_233 = tensor.collapse_shape %317 [[0], [1, 2], [3]] : tensor<1x32x16x4096xf16> into tensor<1x512x4096xf16>
    %expanded_234 = tensor.expand_shape %collapsed_233 [[0], [1], [2, 3]] : tensor<1x512x4096xf16> into tensor<1x512x64x64xf16>
    %expanded_235 = tensor.expand_shape %cst_92 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %expanded_236 = tensor.expand_shape %cst_91 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %318 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_234, %expanded_236 : tensor<1x512x64x64xf16>, tensor<1x512x1x1xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %319 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%318, %expanded_235 : tensor<1x512x64x64xf16>, tensor<1x512x1x1xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %320 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%319 : tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.negf %in : f16
      %884 = math.exp %883 : f16
      %885 = arith.addf %884, %cst_141 : f16
      %886 = arith.divf %cst_141, %885 : f16
      linalg.yield %886 : f16
    } -> tensor<1x512x64x64xf16>
    %321 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%320, %319 : tensor<1x512x64x64xf16>, tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %padded_237 = tensor.pad %321 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x512x64x64xf16> to tensor<1x512x66x66xf16>
    %322 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_89 : tensor<512xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x64x64xf16>
    %323 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_237, %cst_90 : tensor<1x512x66x66xf16>, tensor<512x512x3x3xf16>) outs(%322 : tensor<1x512x64x64xf16>) -> tensor<1x512x64x64xf16>
    %collapsed_238 = tensor.collapse_shape %323 [[0], [1], [2, 3]] : tensor<1x512x64x64xf16> into tensor<1x512x4096xf16>
    %expanded_239 = tensor.expand_shape %collapsed_238 [[0], [1, 2], [3]] : tensor<1x512x4096xf16> into tensor<1x32x16x4096xf16>
    %324 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_239 : tensor<1x32x16x4096xf16>) outs(%7 : tensor<1x32x16x4096xf32>) {
    ^bb0(%in: f16, %out: f32):
      %883 = arith.extf %in : f16 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x16x4096xf32>
    %325 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%324 : tensor<1x32x16x4096xf32>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f32, %out: f64):
      %883 = arith.extf %in : f32 to f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %326 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%325 : tensor<1x32x16x4096xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %327 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%326 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_147 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %328 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%325, %327 : tensor<1x32x16x4096xf64>, tensor<1x32x1x1xf64>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.subf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %329 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%328, %328 : tensor<1x32x16x4096xf64>, tensor<1x32x16x4096xf64>) outs(%9 : tensor<1x32x16x4096xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.mulf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x4096xf64>
    %330 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%329 : tensor<1x32x16x4096xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %331 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%330 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_147 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %332 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%331 : tensor<1x32x1x1xf64>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f64, %out: f32):
      %883 = arith.truncf %in : f64 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %333 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%324 : tensor<1x32x16x4096xf32>) outs(%21 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.addf %in, %out : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %334 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%333 : tensor<1x32x1x1xf32>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.divf %in, %cst_148 : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %335 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%332 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %336 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%335 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %337 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%334 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %338 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%337 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %339 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%336 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_145 : f64 to f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x32x1x1xf16>
    %340 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%339 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.rsqrt %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %341 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_239, %338 : tensor<1x32x16x4096xf16>, tensor<1x32x1x1xf16>) outs(%31 : tensor<1x32x16x4096xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x4096xf16>
    %342 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%341, %340 : tensor<1x32x16x4096xf16>, tensor<1x32x1x1xf16>) outs(%31 : tensor<1x32x16x4096xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x4096xf16>
    %collapsed_240 = tensor.collapse_shape %342 [[0], [1, 2], [3]] : tensor<1x32x16x4096xf16> into tensor<1x512x4096xf16>
    %expanded_241 = tensor.expand_shape %collapsed_240 [[0], [1], [2, 3]] : tensor<1x512x4096xf16> into tensor<1x512x64x64xf16>
    %expanded_242 = tensor.expand_shape %cst_88 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %expanded_243 = tensor.expand_shape %cst_87 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %343 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_241, %expanded_243 : tensor<1x512x64x64xf16>, tensor<1x512x1x1xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %344 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%343, %expanded_242 : tensor<1x512x64x64xf16>, tensor<1x512x1x1xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %345 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%344 : tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.negf %in : f16
      %884 = math.exp %883 : f16
      %885 = arith.addf %884, %cst_141 : f16
      %886 = arith.divf %cst_141, %885 : f16
      linalg.yield %886 : f16
    } -> tensor<1x512x64x64xf16>
    %346 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%345, %344 : tensor<1x512x64x64xf16>, tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %padded_244 = tensor.pad %346 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x512x64x64xf16> to tensor<1x512x66x66xf16>
    %347 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_85 : tensor<512xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x64x64xf16>
    %348 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_244, %cst_86 : tensor<1x512x66x66xf16>, tensor<512x512x3x3xf16>) outs(%347 : tensor<1x512x64x64xf16>) -> tensor<1x512x64x64xf16>
    %349 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%297, %348 : tensor<1x512x64x64xf16>, tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x64x64xf16>
    %350 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%349 : tensor<1x512x64x64xf16>) outs(%4 : tensor<1x512x64x64xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x64x64xf16>
    %351 = tensor.empty() : tensor<1x512x128x128xf16>
    %352 = linalg.generic {indexing_maps = [#map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%out: f16):
      %883 = linalg.index 0 : index
      %884 = linalg.index 1 : index
      %885 = linalg.index 2 : index
      %886 = linalg.index 3 : index
      %887 = arith.floordivsi %885, %c2 : index
      %888 = arith.floordivsi %886, %c2 : index
      %extracted = tensor.extract %350[%883, %884, %887, %888] : tensor<1x512x64x64xf16>
      linalg.yield %extracted : f16
    } -> tensor<1x512x128x128xf16>
    %padded_245 = tensor.pad %352 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x512x128x128xf16> to tensor<1x512x130x130xf16>
    %353 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_83 : tensor<512xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x128x128xf16>
    %354 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_245, %cst_84 : tensor<1x512x130x130xf16>, tensor<512x512x3x3xf16>) outs(%353 : tensor<1x512x128x128xf16>) -> tensor<1x512x128x128xf16>
    %collapsed_246 = tensor.collapse_shape %354 [[0], [1], [2, 3]] : tensor<1x512x128x128xf16> into tensor<1x512x16384xf16>
    %expanded_247 = tensor.expand_shape %collapsed_246 [[0], [1, 2], [3]] : tensor<1x512x16384xf16> into tensor<1x32x16x16384xf16>
    %355 = tensor.empty() : tensor<1x32x16x16384xf32>
    %356 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_247 : tensor<1x32x16x16384xf16>) outs(%355 : tensor<1x32x16x16384xf32>) {
    ^bb0(%in: f16, %out: f32):
      %883 = arith.extf %in : f16 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x16x16384xf32>
    %357 = tensor.empty() : tensor<1x32x16x16384xf64>
    %358 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%356 : tensor<1x32x16x16384xf32>) outs(%357 : tensor<1x32x16x16384xf64>) {
    ^bb0(%in: f32, %out: f64):
      %883 = arith.extf %in : f32 to f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x16384xf64>
    %359 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%358 : tensor<1x32x16x16384xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %360 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%359 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_149 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %361 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%358, %360 : tensor<1x32x16x16384xf64>, tensor<1x32x1x1xf64>) outs(%357 : tensor<1x32x16x16384xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.subf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x16384xf64>
    %362 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%361, %361 : tensor<1x32x16x16384xf64>, tensor<1x32x16x16384xf64>) outs(%357 : tensor<1x32x16x16384xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.mulf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x16384xf64>
    %363 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%362 : tensor<1x32x16x16384xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %364 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%363 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_149 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %365 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%364 : tensor<1x32x1x1xf64>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f64, %out: f32):
      %883 = arith.truncf %in : f64 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %366 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%356 : tensor<1x32x16x16384xf32>) outs(%21 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.addf %in, %out : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %367 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%366 : tensor<1x32x1x1xf32>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.divf %in, %cst_150 : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %368 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%365 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %369 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%368 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %370 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%367 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %371 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%370 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %372 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%369 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_145 : f64 to f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x32x1x1xf16>
    %373 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%372 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.rsqrt %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %374 = tensor.empty() : tensor<1x32x16x16384xf16>
    %375 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_247, %371 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%374 : tensor<1x32x16x16384xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x16384xf16>
    %376 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%375, %373 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%374 : tensor<1x32x16x16384xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x16384xf16>
    %collapsed_248 = tensor.collapse_shape %376 [[0], [1, 2], [3]] : tensor<1x32x16x16384xf16> into tensor<1x512x16384xf16>
    %expanded_249 = tensor.expand_shape %collapsed_248 [[0], [1], [2, 3]] : tensor<1x512x16384xf16> into tensor<1x512x128x128xf16>
    %expanded_250 = tensor.expand_shape %cst_82 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %expanded_251 = tensor.expand_shape %cst_81 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %377 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_249, %expanded_251 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x128x128xf16>
    %378 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%377, %expanded_250 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x128x128xf16>
    %379 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%378 : tensor<1x512x128x128xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.negf %in : f16
      %884 = math.exp %883 : f16
      %885 = arith.addf %884, %cst_141 : f16
      %886 = arith.divf %cst_141, %885 : f16
      linalg.yield %886 : f16
    } -> tensor<1x512x128x128xf16>
    %380 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%379, %378 : tensor<1x512x128x128xf16>, tensor<1x512x128x128xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x128x128xf16>
    %padded_252 = tensor.pad %380 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x512x128x128xf16> to tensor<1x512x130x130xf16>
    %381 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_79 : tensor<512xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x128x128xf16>
    %382 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_252, %cst_80 : tensor<1x512x130x130xf16>, tensor<512x512x3x3xf16>) outs(%381 : tensor<1x512x128x128xf16>) -> tensor<1x512x128x128xf16>
    %collapsed_253 = tensor.collapse_shape %382 [[0], [1], [2, 3]] : tensor<1x512x128x128xf16> into tensor<1x512x16384xf16>
    %expanded_254 = tensor.expand_shape %collapsed_253 [[0], [1, 2], [3]] : tensor<1x512x16384xf16> into tensor<1x32x16x16384xf16>
    %383 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_254 : tensor<1x32x16x16384xf16>) outs(%355 : tensor<1x32x16x16384xf32>) {
    ^bb0(%in: f16, %out: f32):
      %883 = arith.extf %in : f16 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x16x16384xf32>
    %384 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%383 : tensor<1x32x16x16384xf32>) outs(%357 : tensor<1x32x16x16384xf64>) {
    ^bb0(%in: f32, %out: f64):
      %883 = arith.extf %in : f32 to f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x16384xf64>
    %385 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%384 : tensor<1x32x16x16384xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %386 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%385 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_149 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %387 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%384, %386 : tensor<1x32x16x16384xf64>, tensor<1x32x1x1xf64>) outs(%357 : tensor<1x32x16x16384xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.subf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x16384xf64>
    %388 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%387, %387 : tensor<1x32x16x16384xf64>, tensor<1x32x16x16384xf64>) outs(%357 : tensor<1x32x16x16384xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.mulf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x16384xf64>
    %389 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%388 : tensor<1x32x16x16384xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %390 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%389 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_149 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %391 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%390 : tensor<1x32x1x1xf64>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f64, %out: f32):
      %883 = arith.truncf %in : f64 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %392 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%383 : tensor<1x32x16x16384xf32>) outs(%21 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.addf %in, %out : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %393 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%392 : tensor<1x32x1x1xf32>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.divf %in, %cst_150 : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %394 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%391 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %395 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%394 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %396 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%393 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %397 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%396 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %398 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%395 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_145 : f64 to f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x32x1x1xf16>
    %399 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%398 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.rsqrt %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %400 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_254, %397 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%374 : tensor<1x32x16x16384xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x16384xf16>
    %401 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%400, %399 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%374 : tensor<1x32x16x16384xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x16384xf16>
    %collapsed_255 = tensor.collapse_shape %401 [[0], [1, 2], [3]] : tensor<1x32x16x16384xf16> into tensor<1x512x16384xf16>
    %expanded_256 = tensor.expand_shape %collapsed_255 [[0], [1], [2, 3]] : tensor<1x512x16384xf16> into tensor<1x512x128x128xf16>
    %expanded_257 = tensor.expand_shape %cst_78 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %expanded_258 = tensor.expand_shape %cst_77 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %402 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_256, %expanded_258 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x128x128xf16>
    %403 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%402, %expanded_257 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x128x128xf16>
    %404 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%403 : tensor<1x512x128x128xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.negf %in : f16
      %884 = math.exp %883 : f16
      %885 = arith.addf %884, %cst_141 : f16
      %886 = arith.divf %cst_141, %885 : f16
      linalg.yield %886 : f16
    } -> tensor<1x512x128x128xf16>
    %405 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%404, %403 : tensor<1x512x128x128xf16>, tensor<1x512x128x128xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x128x128xf16>
    %padded_259 = tensor.pad %405 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x512x128x128xf16> to tensor<1x512x130x130xf16>
    %406 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_75 : tensor<512xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x128x128xf16>
    %407 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_259, %cst_76 : tensor<1x512x130x130xf16>, tensor<512x512x3x3xf16>) outs(%406 : tensor<1x512x128x128xf16>) -> tensor<1x512x128x128xf16>
    %408 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%354, %407 : tensor<1x512x128x128xf16>, tensor<1x512x128x128xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x128x128xf16>
    %409 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%408 : tensor<1x512x128x128xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x128x128xf16>
    %collapsed_260 = tensor.collapse_shape %409 [[0], [1], [2, 3]] : tensor<1x512x128x128xf16> into tensor<1x512x16384xf16>
    %expanded_261 = tensor.expand_shape %collapsed_260 [[0], [1, 2], [3]] : tensor<1x512x16384xf16> into tensor<1x32x16x16384xf16>
    %410 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_261 : tensor<1x32x16x16384xf16>) outs(%355 : tensor<1x32x16x16384xf32>) {
    ^bb0(%in: f16, %out: f32):
      %883 = arith.extf %in : f16 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x16x16384xf32>
    %411 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%410 : tensor<1x32x16x16384xf32>) outs(%357 : tensor<1x32x16x16384xf64>) {
    ^bb0(%in: f32, %out: f64):
      %883 = arith.extf %in : f32 to f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x16384xf64>
    %412 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%411 : tensor<1x32x16x16384xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %413 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%412 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_149 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %414 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%411, %413 : tensor<1x32x16x16384xf64>, tensor<1x32x1x1xf64>) outs(%357 : tensor<1x32x16x16384xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.subf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x16384xf64>
    %415 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%414, %414 : tensor<1x32x16x16384xf64>, tensor<1x32x16x16384xf64>) outs(%357 : tensor<1x32x16x16384xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.mulf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x16384xf64>
    %416 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%415 : tensor<1x32x16x16384xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %417 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%416 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_149 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %418 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%417 : tensor<1x32x1x1xf64>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f64, %out: f32):
      %883 = arith.truncf %in : f64 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %419 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%410 : tensor<1x32x16x16384xf32>) outs(%21 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.addf %in, %out : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %420 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%419 : tensor<1x32x1x1xf32>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.divf %in, %cst_150 : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %421 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%418 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %422 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%421 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %423 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%420 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %424 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%423 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %425 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%422 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_145 : f64 to f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x32x1x1xf16>
    %426 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%425 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.rsqrt %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %427 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_261, %424 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%374 : tensor<1x32x16x16384xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x16384xf16>
    %428 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%427, %426 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%374 : tensor<1x32x16x16384xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x16384xf16>
    %collapsed_262 = tensor.collapse_shape %428 [[0], [1, 2], [3]] : tensor<1x32x16x16384xf16> into tensor<1x512x16384xf16>
    %expanded_263 = tensor.expand_shape %collapsed_262 [[0], [1], [2, 3]] : tensor<1x512x16384xf16> into tensor<1x512x128x128xf16>
    %expanded_264 = tensor.expand_shape %cst_74 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %expanded_265 = tensor.expand_shape %cst_73 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %429 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_263, %expanded_265 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x128x128xf16>
    %430 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%429, %expanded_264 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x128x128xf16>
    %431 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%430 : tensor<1x512x128x128xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.negf %in : f16
      %884 = math.exp %883 : f16
      %885 = arith.addf %884, %cst_141 : f16
      %886 = arith.divf %cst_141, %885 : f16
      linalg.yield %886 : f16
    } -> tensor<1x512x128x128xf16>
    %432 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%431, %430 : tensor<1x512x128x128xf16>, tensor<1x512x128x128xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x128x128xf16>
    %padded_266 = tensor.pad %432 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x512x128x128xf16> to tensor<1x512x130x130xf16>
    %433 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_71 : tensor<512xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x128x128xf16>
    %434 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_266, %cst_72 : tensor<1x512x130x130xf16>, tensor<512x512x3x3xf16>) outs(%433 : tensor<1x512x128x128xf16>) -> tensor<1x512x128x128xf16>
    %collapsed_267 = tensor.collapse_shape %434 [[0], [1], [2, 3]] : tensor<1x512x128x128xf16> into tensor<1x512x16384xf16>
    %expanded_268 = tensor.expand_shape %collapsed_267 [[0], [1, 2], [3]] : tensor<1x512x16384xf16> into tensor<1x32x16x16384xf16>
    %435 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_268 : tensor<1x32x16x16384xf16>) outs(%355 : tensor<1x32x16x16384xf32>) {
    ^bb0(%in: f16, %out: f32):
      %883 = arith.extf %in : f16 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x16x16384xf32>
    %436 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%435 : tensor<1x32x16x16384xf32>) outs(%357 : tensor<1x32x16x16384xf64>) {
    ^bb0(%in: f32, %out: f64):
      %883 = arith.extf %in : f32 to f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x16384xf64>
    %437 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%436 : tensor<1x32x16x16384xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %438 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%437 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_149 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %439 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%436, %438 : tensor<1x32x16x16384xf64>, tensor<1x32x1x1xf64>) outs(%357 : tensor<1x32x16x16384xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.subf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x16384xf64>
    %440 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%439, %439 : tensor<1x32x16x16384xf64>, tensor<1x32x16x16384xf64>) outs(%357 : tensor<1x32x16x16384xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.mulf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x16384xf64>
    %441 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%440 : tensor<1x32x16x16384xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %442 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%441 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_149 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %443 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%442 : tensor<1x32x1x1xf64>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f64, %out: f32):
      %883 = arith.truncf %in : f64 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %444 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%435 : tensor<1x32x16x16384xf32>) outs(%21 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.addf %in, %out : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %445 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%444 : tensor<1x32x1x1xf32>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.divf %in, %cst_150 : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %446 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%443 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %447 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%446 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %448 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%445 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %449 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%448 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %450 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%447 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_145 : f64 to f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x32x1x1xf16>
    %451 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%450 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.rsqrt %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %452 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_268, %449 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%374 : tensor<1x32x16x16384xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x16384xf16>
    %453 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%452, %451 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%374 : tensor<1x32x16x16384xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x16384xf16>
    %collapsed_269 = tensor.collapse_shape %453 [[0], [1, 2], [3]] : tensor<1x32x16x16384xf16> into tensor<1x512x16384xf16>
    %expanded_270 = tensor.expand_shape %collapsed_269 [[0], [1], [2, 3]] : tensor<1x512x16384xf16> into tensor<1x512x128x128xf16>
    %expanded_271 = tensor.expand_shape %cst_70 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %expanded_272 = tensor.expand_shape %cst_69 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %454 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_270, %expanded_272 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x128x128xf16>
    %455 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%454, %expanded_271 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x128x128xf16>
    %456 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%455 : tensor<1x512x128x128xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.negf %in : f16
      %884 = math.exp %883 : f16
      %885 = arith.addf %884, %cst_141 : f16
      %886 = arith.divf %cst_141, %885 : f16
      linalg.yield %886 : f16
    } -> tensor<1x512x128x128xf16>
    %457 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%456, %455 : tensor<1x512x128x128xf16>, tensor<1x512x128x128xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x128x128xf16>
    %padded_273 = tensor.pad %457 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x512x128x128xf16> to tensor<1x512x130x130xf16>
    %458 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_67 : tensor<512xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x128x128xf16>
    %459 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_273, %cst_68 : tensor<1x512x130x130xf16>, tensor<512x512x3x3xf16>) outs(%458 : tensor<1x512x128x128xf16>) -> tensor<1x512x128x128xf16>
    %460 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%409, %459 : tensor<1x512x128x128xf16>, tensor<1x512x128x128xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x128x128xf16>
    %461 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%460 : tensor<1x512x128x128xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x128x128xf16>
    %collapsed_274 = tensor.collapse_shape %461 [[0], [1], [2, 3]] : tensor<1x512x128x128xf16> into tensor<1x512x16384xf16>
    %expanded_275 = tensor.expand_shape %collapsed_274 [[0], [1, 2], [3]] : tensor<1x512x16384xf16> into tensor<1x32x16x16384xf16>
    %462 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_275 : tensor<1x32x16x16384xf16>) outs(%355 : tensor<1x32x16x16384xf32>) {
    ^bb0(%in: f16, %out: f32):
      %883 = arith.extf %in : f16 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x16x16384xf32>
    %463 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%462 : tensor<1x32x16x16384xf32>) outs(%357 : tensor<1x32x16x16384xf64>) {
    ^bb0(%in: f32, %out: f64):
      %883 = arith.extf %in : f32 to f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x16384xf64>
    %464 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%463 : tensor<1x32x16x16384xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %465 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%464 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_149 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %466 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%463, %465 : tensor<1x32x16x16384xf64>, tensor<1x32x1x1xf64>) outs(%357 : tensor<1x32x16x16384xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.subf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x16384xf64>
    %467 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%466, %466 : tensor<1x32x16x16384xf64>, tensor<1x32x16x16384xf64>) outs(%357 : tensor<1x32x16x16384xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.mulf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x16384xf64>
    %468 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%467 : tensor<1x32x16x16384xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %469 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%468 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_149 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %470 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%469 : tensor<1x32x1x1xf64>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f64, %out: f32):
      %883 = arith.truncf %in : f64 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %471 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%462 : tensor<1x32x16x16384xf32>) outs(%21 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.addf %in, %out : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %472 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%471 : tensor<1x32x1x1xf32>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.divf %in, %cst_150 : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %473 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%470 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %474 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%473 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %475 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%472 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %476 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%475 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %477 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%474 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_145 : f64 to f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x32x1x1xf16>
    %478 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%477 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.rsqrt %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %479 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_275, %476 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%374 : tensor<1x32x16x16384xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x16384xf16>
    %480 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%479, %478 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%374 : tensor<1x32x16x16384xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x16384xf16>
    %collapsed_276 = tensor.collapse_shape %480 [[0], [1, 2], [3]] : tensor<1x32x16x16384xf16> into tensor<1x512x16384xf16>
    %expanded_277 = tensor.expand_shape %collapsed_276 [[0], [1], [2, 3]] : tensor<1x512x16384xf16> into tensor<1x512x128x128xf16>
    %expanded_278 = tensor.expand_shape %cst_66 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %expanded_279 = tensor.expand_shape %cst_65 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %481 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_277, %expanded_279 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x128x128xf16>
    %482 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%481, %expanded_278 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x128x128xf16>
    %483 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%482 : tensor<1x512x128x128xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.negf %in : f16
      %884 = math.exp %883 : f16
      %885 = arith.addf %884, %cst_141 : f16
      %886 = arith.divf %cst_141, %885 : f16
      linalg.yield %886 : f16
    } -> tensor<1x512x128x128xf16>
    %484 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%483, %482 : tensor<1x512x128x128xf16>, tensor<1x512x128x128xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x128x128xf16>
    %padded_280 = tensor.pad %484 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x512x128x128xf16> to tensor<1x512x130x130xf16>
    %485 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_63 : tensor<512xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x128x128xf16>
    %486 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_280, %cst_64 : tensor<1x512x130x130xf16>, tensor<512x512x3x3xf16>) outs(%485 : tensor<1x512x128x128xf16>) -> tensor<1x512x128x128xf16>
    %collapsed_281 = tensor.collapse_shape %486 [[0], [1], [2, 3]] : tensor<1x512x128x128xf16> into tensor<1x512x16384xf16>
    %expanded_282 = tensor.expand_shape %collapsed_281 [[0], [1, 2], [3]] : tensor<1x512x16384xf16> into tensor<1x32x16x16384xf16>
    %487 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_282 : tensor<1x32x16x16384xf16>) outs(%355 : tensor<1x32x16x16384xf32>) {
    ^bb0(%in: f16, %out: f32):
      %883 = arith.extf %in : f16 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x16x16384xf32>
    %488 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%487 : tensor<1x32x16x16384xf32>) outs(%357 : tensor<1x32x16x16384xf64>) {
    ^bb0(%in: f32, %out: f64):
      %883 = arith.extf %in : f32 to f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x16384xf64>
    %489 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%488 : tensor<1x32x16x16384xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %490 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%489 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_149 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %491 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%488, %490 : tensor<1x32x16x16384xf64>, tensor<1x32x1x1xf64>) outs(%357 : tensor<1x32x16x16384xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.subf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x16384xf64>
    %492 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%491, %491 : tensor<1x32x16x16384xf64>, tensor<1x32x16x16384xf64>) outs(%357 : tensor<1x32x16x16384xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.mulf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x16384xf64>
    %493 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%492 : tensor<1x32x16x16384xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %494 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%493 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_149 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %495 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%494 : tensor<1x32x1x1xf64>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f64, %out: f32):
      %883 = arith.truncf %in : f64 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %496 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%487 : tensor<1x32x16x16384xf32>) outs(%21 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.addf %in, %out : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %497 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%496 : tensor<1x32x1x1xf32>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.divf %in, %cst_150 : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %498 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%495 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %499 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%498 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %500 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%497 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %501 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%500 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %502 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%499 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_145 : f64 to f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x32x1x1xf16>
    %503 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%502 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.rsqrt %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %504 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_282, %501 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%374 : tensor<1x32x16x16384xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x16384xf16>
    %505 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%504, %503 : tensor<1x32x16x16384xf16>, tensor<1x32x1x1xf16>) outs(%374 : tensor<1x32x16x16384xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x16384xf16>
    %collapsed_283 = tensor.collapse_shape %505 [[0], [1, 2], [3]] : tensor<1x32x16x16384xf16> into tensor<1x512x16384xf16>
    %expanded_284 = tensor.expand_shape %collapsed_283 [[0], [1], [2, 3]] : tensor<1x512x16384xf16> into tensor<1x512x128x128xf16>
    %expanded_285 = tensor.expand_shape %cst_62 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %expanded_286 = tensor.expand_shape %cst_61 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %506 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_284, %expanded_286 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x128x128xf16>
    %507 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%506, %expanded_285 : tensor<1x512x128x128xf16>, tensor<1x512x1x1xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x128x128xf16>
    %508 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%507 : tensor<1x512x128x128xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.negf %in : f16
      %884 = math.exp %883 : f16
      %885 = arith.addf %884, %cst_141 : f16
      %886 = arith.divf %cst_141, %885 : f16
      linalg.yield %886 : f16
    } -> tensor<1x512x128x128xf16>
    %509 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%508, %507 : tensor<1x512x128x128xf16>, tensor<1x512x128x128xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x128x128xf16>
    %padded_287 = tensor.pad %509 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x512x128x128xf16> to tensor<1x512x130x130xf16>
    %510 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_59 : tensor<512xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x128x128xf16>
    %511 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_287, %cst_60 : tensor<1x512x130x130xf16>, tensor<512x512x3x3xf16>) outs(%510 : tensor<1x512x128x128xf16>) -> tensor<1x512x128x128xf16>
    %512 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%461, %511 : tensor<1x512x128x128xf16>, tensor<1x512x128x128xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x128x128xf16>
    %513 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%512 : tensor<1x512x128x128xf16>) outs(%351 : tensor<1x512x128x128xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x128x128xf16>
    %514 = tensor.empty() : tensor<1x512x256x256xf16>
    %515 = linalg.generic {indexing_maps = [#map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} outs(%514 : tensor<1x512x256x256xf16>) {
    ^bb0(%out: f16):
      %883 = linalg.index 0 : index
      %884 = linalg.index 1 : index
      %885 = linalg.index 2 : index
      %886 = linalg.index 3 : index
      %887 = arith.floordivsi %885, %c2 : index
      %888 = arith.floordivsi %886, %c2 : index
      %extracted = tensor.extract %513[%883, %884, %887, %888] : tensor<1x512x128x128xf16>
      linalg.yield %extracted : f16
    } -> tensor<1x512x256x256xf16>
    %padded_288 = tensor.pad %515 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x512x256x256xf16> to tensor<1x512x258x258xf16>
    %516 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_57 : tensor<512xf16>) outs(%514 : tensor<1x512x256x256xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x512x256x256xf16>
    %517 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_288, %cst_58 : tensor<1x512x258x258xf16>, tensor<512x512x3x3xf16>) outs(%516 : tensor<1x512x256x256xf16>) -> tensor<1x512x256x256xf16>
    %collapsed_289 = tensor.collapse_shape %517 [[0], [1], [2, 3]] : tensor<1x512x256x256xf16> into tensor<1x512x65536xf16>
    %expanded_290 = tensor.expand_shape %collapsed_289 [[0], [1, 2], [3]] : tensor<1x512x65536xf16> into tensor<1x32x16x65536xf16>
    %518 = tensor.empty() : tensor<1x32x16x65536xf32>
    %519 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_290 : tensor<1x32x16x65536xf16>) outs(%518 : tensor<1x32x16x65536xf32>) {
    ^bb0(%in: f16, %out: f32):
      %883 = arith.extf %in : f16 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x16x65536xf32>
    %520 = tensor.empty() : tensor<1x32x16x65536xf64>
    %521 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%519 : tensor<1x32x16x65536xf32>) outs(%520 : tensor<1x32x16x65536xf64>) {
    ^bb0(%in: f32, %out: f64):
      %883 = arith.extf %in : f32 to f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x65536xf64>
    %522 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%521 : tensor<1x32x16x65536xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %523 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%522 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_151 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %524 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%521, %523 : tensor<1x32x16x65536xf64>, tensor<1x32x1x1xf64>) outs(%520 : tensor<1x32x16x65536xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.subf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x65536xf64>
    %525 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%524, %524 : tensor<1x32x16x65536xf64>, tensor<1x32x16x65536xf64>) outs(%520 : tensor<1x32x16x65536xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.mulf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x16x65536xf64>
    %526 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%525 : tensor<1x32x16x65536xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %527 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%526 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_151 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %528 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%527 : tensor<1x32x1x1xf64>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f64, %out: f32):
      %883 = arith.truncf %in : f64 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %529 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%519 : tensor<1x32x16x65536xf32>) outs(%21 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.addf %in, %out : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %530 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%529 : tensor<1x32x1x1xf32>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.divf %in, %cst_152 : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %531 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%528 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %532 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%531 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %533 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%530 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %534 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%533 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %535 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%532 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_145 : f64 to f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x32x1x1xf16>
    %536 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%535 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.rsqrt %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %537 = tensor.empty() : tensor<1x32x16x65536xf16>
    %538 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_290, %534 : tensor<1x32x16x65536xf16>, tensor<1x32x1x1xf16>) outs(%537 : tensor<1x32x16x65536xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x65536xf16>
    %539 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%538, %536 : tensor<1x32x16x65536xf16>, tensor<1x32x1x1xf16>) outs(%537 : tensor<1x32x16x65536xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x16x65536xf16>
    %collapsed_291 = tensor.collapse_shape %539 [[0], [1, 2], [3]] : tensor<1x32x16x65536xf16> into tensor<1x512x65536xf16>
    %expanded_292 = tensor.expand_shape %collapsed_291 [[0], [1], [2, 3]] : tensor<1x512x65536xf16> into tensor<1x512x256x256xf16>
    %expanded_293 = tensor.expand_shape %cst_56 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %expanded_294 = tensor.expand_shape %cst_55 [[0, 1, 2, 3]] : tensor<512xf16> into tensor<1x512x1x1xf16>
    %540 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_292, %expanded_294 : tensor<1x512x256x256xf16>, tensor<1x512x1x1xf16>) outs(%514 : tensor<1x512x256x256xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x256x256xf16>
    %541 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%540, %expanded_293 : tensor<1x512x256x256xf16>, tensor<1x512x1x1xf16>) outs(%514 : tensor<1x512x256x256xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x256x256xf16>
    %542 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%541 : tensor<1x512x256x256xf16>) outs(%514 : tensor<1x512x256x256xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.negf %in : f16
      %884 = math.exp %883 : f16
      %885 = arith.addf %884, %cst_141 : f16
      %886 = arith.divf %cst_141, %885 : f16
      linalg.yield %886 : f16
    } -> tensor<1x512x256x256xf16>
    %543 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%542, %541 : tensor<1x512x256x256xf16>, tensor<1x512x256x256xf16>) outs(%514 : tensor<1x512x256x256xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x512x256x256xf16>
    %padded_295 = tensor.pad %543 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x512x256x256xf16> to tensor<1x512x258x258xf16>
    %544 = tensor.empty() : tensor<1x256x256x256xf16>
    %545 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_53 : tensor<256xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x256x256x256xf16>
    %546 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_295, %cst_54 : tensor<1x512x258x258xf16>, tensor<256x512x3x3xf16>) outs(%545 : tensor<1x256x256x256xf16>) -> tensor<1x256x256x256xf16>
    %collapsed_296 = tensor.collapse_shape %546 [[0], [1], [2, 3]] : tensor<1x256x256x256xf16> into tensor<1x256x65536xf16>
    %expanded_297 = tensor.expand_shape %collapsed_296 [[0], [1, 2], [3]] : tensor<1x256x65536xf16> into tensor<1x32x8x65536xf16>
    %547 = tensor.empty() : tensor<1x32x8x65536xf32>
    %548 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_297 : tensor<1x32x8x65536xf16>) outs(%547 : tensor<1x32x8x65536xf32>) {
    ^bb0(%in: f16, %out: f32):
      %883 = arith.extf %in : f16 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x8x65536xf32>
    %549 = tensor.empty() : tensor<1x32x8x65536xf64>
    %550 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%548 : tensor<1x32x8x65536xf32>) outs(%549 : tensor<1x32x8x65536xf64>) {
    ^bb0(%in: f32, %out: f64):
      %883 = arith.extf %in : f32 to f64
      linalg.yield %883 : f64
    } -> tensor<1x32x8x65536xf64>
    %551 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%550 : tensor<1x32x8x65536xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %552 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%551 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_153 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %553 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%550, %552 : tensor<1x32x8x65536xf64>, tensor<1x32x1x1xf64>) outs(%549 : tensor<1x32x8x65536xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.subf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x8x65536xf64>
    %554 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%553, %553 : tensor<1x32x8x65536xf64>, tensor<1x32x8x65536xf64>) outs(%549 : tensor<1x32x8x65536xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.mulf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x8x65536xf64>
    %555 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%554 : tensor<1x32x8x65536xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %556 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%555 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_153 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %557 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%556 : tensor<1x32x1x1xf64>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f64, %out: f32):
      %883 = arith.truncf %in : f64 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %558 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%548 : tensor<1x32x8x65536xf32>) outs(%21 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.addf %in, %out : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %559 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%558 : tensor<1x32x1x1xf32>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.divf %in, %cst_154 : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %560 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%557 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %561 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%560 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %562 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%559 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %563 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%562 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %564 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%561 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_145 : f64 to f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x32x1x1xf16>
    %565 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%564 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.rsqrt %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %566 = tensor.empty() : tensor<1x32x8x65536xf16>
    %567 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_297, %563 : tensor<1x32x8x65536xf16>, tensor<1x32x1x1xf16>) outs(%566 : tensor<1x32x8x65536xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x8x65536xf16>
    %568 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%567, %565 : tensor<1x32x8x65536xf16>, tensor<1x32x1x1xf16>) outs(%566 : tensor<1x32x8x65536xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x8x65536xf16>
    %collapsed_298 = tensor.collapse_shape %568 [[0], [1, 2], [3]] : tensor<1x32x8x65536xf16> into tensor<1x256x65536xf16>
    %expanded_299 = tensor.expand_shape %collapsed_298 [[0], [1], [2, 3]] : tensor<1x256x65536xf16> into tensor<1x256x256x256xf16>
    %expanded_300 = tensor.expand_shape %cst_52 [[0, 1, 2, 3]] : tensor<256xf16> into tensor<1x256x1x1xf16>
    %expanded_301 = tensor.expand_shape %cst_51 [[0, 1, 2, 3]] : tensor<256xf16> into tensor<1x256x1x1xf16>
    %569 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_299, %expanded_301 : tensor<1x256x256x256xf16>, tensor<1x256x1x1xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x256x256x256xf16>
    %570 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%569, %expanded_300 : tensor<1x256x256x256xf16>, tensor<1x256x1x1xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x256x256x256xf16>
    %571 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%570 : tensor<1x256x256x256xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.negf %in : f16
      %884 = math.exp %883 : f16
      %885 = arith.addf %884, %cst_141 : f16
      %886 = arith.divf %cst_141, %885 : f16
      linalg.yield %886 : f16
    } -> tensor<1x256x256x256xf16>
    %572 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%571, %570 : tensor<1x256x256x256xf16>, tensor<1x256x256x256xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x256x256x256xf16>
    %padded_302 = tensor.pad %572 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x256x256x256xf16> to tensor<1x256x258x258xf16>
    %573 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_49 : tensor<256xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x256x256x256xf16>
    %574 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_302, %cst_50 : tensor<1x256x258x258xf16>, tensor<256x256x3x3xf16>) outs(%573 : tensor<1x256x256x256xf16>) -> tensor<1x256x256x256xf16>
    %575 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_47 : tensor<256xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x256x256x256xf16>
    %576 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%517, %cst_48 : tensor<1x512x256x256xf16>, tensor<256x512x1x1xf16>) outs(%575 : tensor<1x256x256x256xf16>) -> tensor<1x256x256x256xf16>
    %577 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%576, %574 : tensor<1x256x256x256xf16>, tensor<1x256x256x256xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x256x256x256xf16>
    %578 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%577 : tensor<1x256x256x256xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x256x256x256xf16>
    %collapsed_303 = tensor.collapse_shape %578 [[0], [1], [2, 3]] : tensor<1x256x256x256xf16> into tensor<1x256x65536xf16>
    %expanded_304 = tensor.expand_shape %collapsed_303 [[0], [1, 2], [3]] : tensor<1x256x65536xf16> into tensor<1x32x8x65536xf16>
    %579 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_304 : tensor<1x32x8x65536xf16>) outs(%547 : tensor<1x32x8x65536xf32>) {
    ^bb0(%in: f16, %out: f32):
      %883 = arith.extf %in : f16 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x8x65536xf32>
    %580 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%579 : tensor<1x32x8x65536xf32>) outs(%549 : tensor<1x32x8x65536xf64>) {
    ^bb0(%in: f32, %out: f64):
      %883 = arith.extf %in : f32 to f64
      linalg.yield %883 : f64
    } -> tensor<1x32x8x65536xf64>
    %581 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%580 : tensor<1x32x8x65536xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %582 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%581 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_153 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %583 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%580, %582 : tensor<1x32x8x65536xf64>, tensor<1x32x1x1xf64>) outs(%549 : tensor<1x32x8x65536xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.subf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x8x65536xf64>
    %584 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%583, %583 : tensor<1x32x8x65536xf64>, tensor<1x32x8x65536xf64>) outs(%549 : tensor<1x32x8x65536xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.mulf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x8x65536xf64>
    %585 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%584 : tensor<1x32x8x65536xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %586 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%585 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_153 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %587 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%586 : tensor<1x32x1x1xf64>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f64, %out: f32):
      %883 = arith.truncf %in : f64 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %588 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%579 : tensor<1x32x8x65536xf32>) outs(%21 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.addf %in, %out : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %589 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%588 : tensor<1x32x1x1xf32>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.divf %in, %cst_154 : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %590 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%587 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %591 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%590 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %592 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%589 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %593 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%592 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %594 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%591 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_145 : f64 to f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x32x1x1xf16>
    %595 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%594 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.rsqrt %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %596 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_304, %593 : tensor<1x32x8x65536xf16>, tensor<1x32x1x1xf16>) outs(%566 : tensor<1x32x8x65536xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x8x65536xf16>
    %597 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%596, %595 : tensor<1x32x8x65536xf16>, tensor<1x32x1x1xf16>) outs(%566 : tensor<1x32x8x65536xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x8x65536xf16>
    %collapsed_305 = tensor.collapse_shape %597 [[0], [1, 2], [3]] : tensor<1x32x8x65536xf16> into tensor<1x256x65536xf16>
    %expanded_306 = tensor.expand_shape %collapsed_305 [[0], [1], [2, 3]] : tensor<1x256x65536xf16> into tensor<1x256x256x256xf16>
    %expanded_307 = tensor.expand_shape %cst_46 [[0, 1, 2, 3]] : tensor<256xf16> into tensor<1x256x1x1xf16>
    %expanded_308 = tensor.expand_shape %cst_45 [[0, 1, 2, 3]] : tensor<256xf16> into tensor<1x256x1x1xf16>
    %598 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_306, %expanded_308 : tensor<1x256x256x256xf16>, tensor<1x256x1x1xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x256x256x256xf16>
    %599 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%598, %expanded_307 : tensor<1x256x256x256xf16>, tensor<1x256x1x1xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x256x256x256xf16>
    %600 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%599 : tensor<1x256x256x256xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.negf %in : f16
      %884 = math.exp %883 : f16
      %885 = arith.addf %884, %cst_141 : f16
      %886 = arith.divf %cst_141, %885 : f16
      linalg.yield %886 : f16
    } -> tensor<1x256x256x256xf16>
    %601 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%600, %599 : tensor<1x256x256x256xf16>, tensor<1x256x256x256xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x256x256x256xf16>
    %padded_309 = tensor.pad %601 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x256x256x256xf16> to tensor<1x256x258x258xf16>
    %602 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_43 : tensor<256xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x256x256x256xf16>
    %603 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_309, %cst_44 : tensor<1x256x258x258xf16>, tensor<256x256x3x3xf16>) outs(%602 : tensor<1x256x256x256xf16>) -> tensor<1x256x256x256xf16>
    %collapsed_310 = tensor.collapse_shape %603 [[0], [1], [2, 3]] : tensor<1x256x256x256xf16> into tensor<1x256x65536xf16>
    %expanded_311 = tensor.expand_shape %collapsed_310 [[0], [1, 2], [3]] : tensor<1x256x65536xf16> into tensor<1x32x8x65536xf16>
    %604 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_311 : tensor<1x32x8x65536xf16>) outs(%547 : tensor<1x32x8x65536xf32>) {
    ^bb0(%in: f16, %out: f32):
      %883 = arith.extf %in : f16 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x8x65536xf32>
    %605 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%604 : tensor<1x32x8x65536xf32>) outs(%549 : tensor<1x32x8x65536xf64>) {
    ^bb0(%in: f32, %out: f64):
      %883 = arith.extf %in : f32 to f64
      linalg.yield %883 : f64
    } -> tensor<1x32x8x65536xf64>
    %606 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%605 : tensor<1x32x8x65536xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %607 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%606 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_153 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %608 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%605, %607 : tensor<1x32x8x65536xf64>, tensor<1x32x1x1xf64>) outs(%549 : tensor<1x32x8x65536xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.subf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x8x65536xf64>
    %609 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%608, %608 : tensor<1x32x8x65536xf64>, tensor<1x32x8x65536xf64>) outs(%549 : tensor<1x32x8x65536xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.mulf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x8x65536xf64>
    %610 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%609 : tensor<1x32x8x65536xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %611 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%610 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_153 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %612 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%611 : tensor<1x32x1x1xf64>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f64, %out: f32):
      %883 = arith.truncf %in : f64 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %613 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%604 : tensor<1x32x8x65536xf32>) outs(%21 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.addf %in, %out : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %614 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%613 : tensor<1x32x1x1xf32>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.divf %in, %cst_154 : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %615 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%612 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %616 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%615 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %617 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%614 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %618 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%617 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %619 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%616 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_145 : f64 to f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x32x1x1xf16>
    %620 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%619 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.rsqrt %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %621 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_311, %618 : tensor<1x32x8x65536xf16>, tensor<1x32x1x1xf16>) outs(%566 : tensor<1x32x8x65536xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x8x65536xf16>
    %622 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%621, %620 : tensor<1x32x8x65536xf16>, tensor<1x32x1x1xf16>) outs(%566 : tensor<1x32x8x65536xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x8x65536xf16>
    %collapsed_312 = tensor.collapse_shape %622 [[0], [1, 2], [3]] : tensor<1x32x8x65536xf16> into tensor<1x256x65536xf16>
    %expanded_313 = tensor.expand_shape %collapsed_312 [[0], [1], [2, 3]] : tensor<1x256x65536xf16> into tensor<1x256x256x256xf16>
    %expanded_314 = tensor.expand_shape %cst_42 [[0, 1, 2, 3]] : tensor<256xf16> into tensor<1x256x1x1xf16>
    %expanded_315 = tensor.expand_shape %cst_41 [[0, 1, 2, 3]] : tensor<256xf16> into tensor<1x256x1x1xf16>
    %623 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_313, %expanded_315 : tensor<1x256x256x256xf16>, tensor<1x256x1x1xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x256x256x256xf16>
    %624 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%623, %expanded_314 : tensor<1x256x256x256xf16>, tensor<1x256x1x1xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x256x256x256xf16>
    %625 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%624 : tensor<1x256x256x256xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.negf %in : f16
      %884 = math.exp %883 : f16
      %885 = arith.addf %884, %cst_141 : f16
      %886 = arith.divf %cst_141, %885 : f16
      linalg.yield %886 : f16
    } -> tensor<1x256x256x256xf16>
    %626 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%625, %624 : tensor<1x256x256x256xf16>, tensor<1x256x256x256xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x256x256x256xf16>
    %padded_316 = tensor.pad %626 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x256x256x256xf16> to tensor<1x256x258x258xf16>
    %627 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_39 : tensor<256xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x256x256x256xf16>
    %628 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_316, %cst_40 : tensor<1x256x258x258xf16>, tensor<256x256x3x3xf16>) outs(%627 : tensor<1x256x256x256xf16>) -> tensor<1x256x256x256xf16>
    %629 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%578, %628 : tensor<1x256x256x256xf16>, tensor<1x256x256x256xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x256x256x256xf16>
    %630 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%629 : tensor<1x256x256x256xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x256x256x256xf16>
    %collapsed_317 = tensor.collapse_shape %630 [[0], [1], [2, 3]] : tensor<1x256x256x256xf16> into tensor<1x256x65536xf16>
    %expanded_318 = tensor.expand_shape %collapsed_317 [[0], [1, 2], [3]] : tensor<1x256x65536xf16> into tensor<1x32x8x65536xf16>
    %631 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_318 : tensor<1x32x8x65536xf16>) outs(%547 : tensor<1x32x8x65536xf32>) {
    ^bb0(%in: f16, %out: f32):
      %883 = arith.extf %in : f16 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x8x65536xf32>
    %632 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%631 : tensor<1x32x8x65536xf32>) outs(%549 : tensor<1x32x8x65536xf64>) {
    ^bb0(%in: f32, %out: f64):
      %883 = arith.extf %in : f32 to f64
      linalg.yield %883 : f64
    } -> tensor<1x32x8x65536xf64>
    %633 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%632 : tensor<1x32x8x65536xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %634 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%633 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_153 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %635 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%632, %634 : tensor<1x32x8x65536xf64>, tensor<1x32x1x1xf64>) outs(%549 : tensor<1x32x8x65536xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.subf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x8x65536xf64>
    %636 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%635, %635 : tensor<1x32x8x65536xf64>, tensor<1x32x8x65536xf64>) outs(%549 : tensor<1x32x8x65536xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.mulf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x8x65536xf64>
    %637 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%636 : tensor<1x32x8x65536xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %638 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%637 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_153 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %639 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%638 : tensor<1x32x1x1xf64>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f64, %out: f32):
      %883 = arith.truncf %in : f64 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %640 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%631 : tensor<1x32x8x65536xf32>) outs(%21 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.addf %in, %out : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %641 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%640 : tensor<1x32x1x1xf32>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.divf %in, %cst_154 : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %642 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%639 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %643 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%642 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %644 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%641 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %645 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%644 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %646 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%643 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_145 : f64 to f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x32x1x1xf16>
    %647 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%646 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.rsqrt %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %648 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_318, %645 : tensor<1x32x8x65536xf16>, tensor<1x32x1x1xf16>) outs(%566 : tensor<1x32x8x65536xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x8x65536xf16>
    %649 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%648, %647 : tensor<1x32x8x65536xf16>, tensor<1x32x1x1xf16>) outs(%566 : tensor<1x32x8x65536xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x8x65536xf16>
    %collapsed_319 = tensor.collapse_shape %649 [[0], [1, 2], [3]] : tensor<1x32x8x65536xf16> into tensor<1x256x65536xf16>
    %expanded_320 = tensor.expand_shape %collapsed_319 [[0], [1], [2, 3]] : tensor<1x256x65536xf16> into tensor<1x256x256x256xf16>
    %expanded_321 = tensor.expand_shape %cst_38 [[0, 1, 2, 3]] : tensor<256xf16> into tensor<1x256x1x1xf16>
    %expanded_322 = tensor.expand_shape %cst_37 [[0, 1, 2, 3]] : tensor<256xf16> into tensor<1x256x1x1xf16>
    %650 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_320, %expanded_322 : tensor<1x256x256x256xf16>, tensor<1x256x1x1xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x256x256x256xf16>
    %651 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%650, %expanded_321 : tensor<1x256x256x256xf16>, tensor<1x256x1x1xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x256x256x256xf16>
    %652 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%651 : tensor<1x256x256x256xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.negf %in : f16
      %884 = math.exp %883 : f16
      %885 = arith.addf %884, %cst_141 : f16
      %886 = arith.divf %cst_141, %885 : f16
      linalg.yield %886 : f16
    } -> tensor<1x256x256x256xf16>
    %653 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%652, %651 : tensor<1x256x256x256xf16>, tensor<1x256x256x256xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x256x256x256xf16>
    %padded_323 = tensor.pad %653 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x256x256x256xf16> to tensor<1x256x258x258xf16>
    %654 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_35 : tensor<256xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x256x256x256xf16>
    %655 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_323, %cst_36 : tensor<1x256x258x258xf16>, tensor<256x256x3x3xf16>) outs(%654 : tensor<1x256x256x256xf16>) -> tensor<1x256x256x256xf16>
    %collapsed_324 = tensor.collapse_shape %655 [[0], [1], [2, 3]] : tensor<1x256x256x256xf16> into tensor<1x256x65536xf16>
    %expanded_325 = tensor.expand_shape %collapsed_324 [[0], [1, 2], [3]] : tensor<1x256x65536xf16> into tensor<1x32x8x65536xf16>
    %656 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_325 : tensor<1x32x8x65536xf16>) outs(%547 : tensor<1x32x8x65536xf32>) {
    ^bb0(%in: f16, %out: f32):
      %883 = arith.extf %in : f16 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x8x65536xf32>
    %657 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%656 : tensor<1x32x8x65536xf32>) outs(%549 : tensor<1x32x8x65536xf64>) {
    ^bb0(%in: f32, %out: f64):
      %883 = arith.extf %in : f32 to f64
      linalg.yield %883 : f64
    } -> tensor<1x32x8x65536xf64>
    %658 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%657 : tensor<1x32x8x65536xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %659 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%658 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_153 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %660 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%657, %659 : tensor<1x32x8x65536xf64>, tensor<1x32x1x1xf64>) outs(%549 : tensor<1x32x8x65536xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.subf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x8x65536xf64>
    %661 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%660, %660 : tensor<1x32x8x65536xf64>, tensor<1x32x8x65536xf64>) outs(%549 : tensor<1x32x8x65536xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.mulf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x8x65536xf64>
    %662 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%661 : tensor<1x32x8x65536xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %663 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%662 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_153 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %664 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%663 : tensor<1x32x1x1xf64>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f64, %out: f32):
      %883 = arith.truncf %in : f64 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %665 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%656 : tensor<1x32x8x65536xf32>) outs(%21 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.addf %in, %out : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %666 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%665 : tensor<1x32x1x1xf32>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.divf %in, %cst_154 : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %667 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%664 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %668 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%667 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %669 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%666 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %670 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%669 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %671 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%668 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_145 : f64 to f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x32x1x1xf16>
    %672 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%671 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.rsqrt %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %673 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_325, %670 : tensor<1x32x8x65536xf16>, tensor<1x32x1x1xf16>) outs(%566 : tensor<1x32x8x65536xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x8x65536xf16>
    %674 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%673, %672 : tensor<1x32x8x65536xf16>, tensor<1x32x1x1xf16>) outs(%566 : tensor<1x32x8x65536xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x8x65536xf16>
    %collapsed_326 = tensor.collapse_shape %674 [[0], [1, 2], [3]] : tensor<1x32x8x65536xf16> into tensor<1x256x65536xf16>
    %expanded_327 = tensor.expand_shape %collapsed_326 [[0], [1], [2, 3]] : tensor<1x256x65536xf16> into tensor<1x256x256x256xf16>
    %expanded_328 = tensor.expand_shape %cst_34 [[0, 1, 2, 3]] : tensor<256xf16> into tensor<1x256x1x1xf16>
    %expanded_329 = tensor.expand_shape %cst_33 [[0, 1, 2, 3]] : tensor<256xf16> into tensor<1x256x1x1xf16>
    %675 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_327, %expanded_329 : tensor<1x256x256x256xf16>, tensor<1x256x1x1xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x256x256x256xf16>
    %676 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%675, %expanded_328 : tensor<1x256x256x256xf16>, tensor<1x256x1x1xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x256x256x256xf16>
    %677 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%676 : tensor<1x256x256x256xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.negf %in : f16
      %884 = math.exp %883 : f16
      %885 = arith.addf %884, %cst_141 : f16
      %886 = arith.divf %cst_141, %885 : f16
      linalg.yield %886 : f16
    } -> tensor<1x256x256x256xf16>
    %678 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%677, %676 : tensor<1x256x256x256xf16>, tensor<1x256x256x256xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x256x256x256xf16>
    %padded_330 = tensor.pad %678 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x256x256x256xf16> to tensor<1x256x258x258xf16>
    %679 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_31 : tensor<256xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x256x256x256xf16>
    %680 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_330, %cst_32 : tensor<1x256x258x258xf16>, tensor<256x256x3x3xf16>) outs(%679 : tensor<1x256x256x256xf16>) -> tensor<1x256x256x256xf16>
    %681 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%630, %680 : tensor<1x256x256x256xf16>, tensor<1x256x256x256xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x256x256x256xf16>
    %682 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%681 : tensor<1x256x256x256xf16>) outs(%544 : tensor<1x256x256x256xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x256x256x256xf16>
    %683 = tensor.empty() : tensor<1x256x512x512xf16>
    %684 = linalg.generic {indexing_maps = [#map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} outs(%683 : tensor<1x256x512x512xf16>) {
    ^bb0(%out: f16):
      %883 = linalg.index 0 : index
      %884 = linalg.index 1 : index
      %885 = linalg.index 2 : index
      %886 = linalg.index 3 : index
      %887 = arith.floordivsi %885, %c2 : index
      %888 = arith.floordivsi %886, %c2 : index
      %extracted = tensor.extract %682[%883, %884, %887, %888] : tensor<1x256x256x256xf16>
      linalg.yield %extracted : f16
    } -> tensor<1x256x512x512xf16>
    %padded_331 = tensor.pad %684 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x256x512x512xf16> to tensor<1x256x514x514xf16>
    %685 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_29 : tensor<256xf16>) outs(%683 : tensor<1x256x512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x256x512x512xf16>
    %686 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_331, %cst_30 : tensor<1x256x514x514xf16>, tensor<256x256x3x3xf16>) outs(%685 : tensor<1x256x512x512xf16>) -> tensor<1x256x512x512xf16>
    %collapsed_332 = tensor.collapse_shape %686 [[0], [1], [2, 3]] : tensor<1x256x512x512xf16> into tensor<1x256x262144xf16>
    %expanded_333 = tensor.expand_shape %collapsed_332 [[0], [1, 2], [3]] : tensor<1x256x262144xf16> into tensor<1x32x8x262144xf16>
    %687 = tensor.empty() : tensor<1x32x8x262144xf32>
    %688 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_333 : tensor<1x32x8x262144xf16>) outs(%687 : tensor<1x32x8x262144xf32>) {
    ^bb0(%in: f16, %out: f32):
      %883 = arith.extf %in : f16 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x8x262144xf32>
    %689 = tensor.empty() : tensor<1x32x8x262144xf64>
    %690 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%688 : tensor<1x32x8x262144xf32>) outs(%689 : tensor<1x32x8x262144xf64>) {
    ^bb0(%in: f32, %out: f64):
      %883 = arith.extf %in : f32 to f64
      linalg.yield %883 : f64
    } -> tensor<1x32x8x262144xf64>
    %691 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%690 : tensor<1x32x8x262144xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %692 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%691 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_155 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %693 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%690, %692 : tensor<1x32x8x262144xf64>, tensor<1x32x1x1xf64>) outs(%689 : tensor<1x32x8x262144xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.subf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x8x262144xf64>
    %694 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%693, %693 : tensor<1x32x8x262144xf64>, tensor<1x32x8x262144xf64>) outs(%689 : tensor<1x32x8x262144xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.mulf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x8x262144xf64>
    %695 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%694 : tensor<1x32x8x262144xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %696 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%695 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_155 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %697 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%696 : tensor<1x32x1x1xf64>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f64, %out: f32):
      %883 = arith.truncf %in : f64 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %698 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%688 : tensor<1x32x8x262144xf32>) outs(%21 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.addf %in, %out : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %699 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%698 : tensor<1x32x1x1xf32>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.divf %in, %cst_156 : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %700 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%697 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %701 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%700 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %702 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%699 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %703 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%702 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %704 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%701 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_145 : f64 to f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x32x1x1xf16>
    %705 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%704 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.rsqrt %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %706 = tensor.empty() : tensor<1x32x8x262144xf16>
    %707 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_333, %703 : tensor<1x32x8x262144xf16>, tensor<1x32x1x1xf16>) outs(%706 : tensor<1x32x8x262144xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x8x262144xf16>
    %708 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%707, %705 : tensor<1x32x8x262144xf16>, tensor<1x32x1x1xf16>) outs(%706 : tensor<1x32x8x262144xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x8x262144xf16>
    %collapsed_334 = tensor.collapse_shape %708 [[0], [1, 2], [3]] : tensor<1x32x8x262144xf16> into tensor<1x256x262144xf16>
    %expanded_335 = tensor.expand_shape %collapsed_334 [[0], [1], [2, 3]] : tensor<1x256x262144xf16> into tensor<1x256x512x512xf16>
    %expanded_336 = tensor.expand_shape %cst_28 [[0, 1, 2, 3]] : tensor<256xf16> into tensor<1x256x1x1xf16>
    %expanded_337 = tensor.expand_shape %cst_27 [[0, 1, 2, 3]] : tensor<256xf16> into tensor<1x256x1x1xf16>
    %709 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_335, %expanded_337 : tensor<1x256x512x512xf16>, tensor<1x256x1x1xf16>) outs(%683 : tensor<1x256x512x512xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x256x512x512xf16>
    %710 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%709, %expanded_336 : tensor<1x256x512x512xf16>, tensor<1x256x1x1xf16>) outs(%683 : tensor<1x256x512x512xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x256x512x512xf16>
    %711 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%710 : tensor<1x256x512x512xf16>) outs(%683 : tensor<1x256x512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.negf %in : f16
      %884 = math.exp %883 : f16
      %885 = arith.addf %884, %cst_141 : f16
      %886 = arith.divf %cst_141, %885 : f16
      linalg.yield %886 : f16
    } -> tensor<1x256x512x512xf16>
    %712 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%711, %710 : tensor<1x256x512x512xf16>, tensor<1x256x512x512xf16>) outs(%683 : tensor<1x256x512x512xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x256x512x512xf16>
    %padded_338 = tensor.pad %712 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x256x512x512xf16> to tensor<1x256x514x514xf16>
    %713 = tensor.empty() : tensor<1x128x512x512xf16>
    %714 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_25 : tensor<128xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x128x512x512xf16>
    %715 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_338, %cst_26 : tensor<1x256x514x514xf16>, tensor<128x256x3x3xf16>) outs(%714 : tensor<1x128x512x512xf16>) -> tensor<1x128x512x512xf16>
    %collapsed_339 = tensor.collapse_shape %715 [[0], [1], [2, 3]] : tensor<1x128x512x512xf16> into tensor<1x128x262144xf16>
    %expanded_340 = tensor.expand_shape %collapsed_339 [[0], [1, 2], [3]] : tensor<1x128x262144xf16> into tensor<1x32x4x262144xf16>
    %716 = tensor.empty() : tensor<1x32x4x262144xf32>
    %717 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_340 : tensor<1x32x4x262144xf16>) outs(%716 : tensor<1x32x4x262144xf32>) {
    ^bb0(%in: f16, %out: f32):
      %883 = arith.extf %in : f16 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x4x262144xf32>
    %718 = tensor.empty() : tensor<1x32x4x262144xf64>
    %719 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%717 : tensor<1x32x4x262144xf32>) outs(%718 : tensor<1x32x4x262144xf64>) {
    ^bb0(%in: f32, %out: f64):
      %883 = arith.extf %in : f32 to f64
      linalg.yield %883 : f64
    } -> tensor<1x32x4x262144xf64>
    %720 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%719 : tensor<1x32x4x262144xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %721 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%720 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_151 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %722 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%719, %721 : tensor<1x32x4x262144xf64>, tensor<1x32x1x1xf64>) outs(%718 : tensor<1x32x4x262144xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.subf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x4x262144xf64>
    %723 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%722, %722 : tensor<1x32x4x262144xf64>, tensor<1x32x4x262144xf64>) outs(%718 : tensor<1x32x4x262144xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.mulf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x4x262144xf64>
    %724 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%723 : tensor<1x32x4x262144xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %725 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%724 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_151 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %726 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%725 : tensor<1x32x1x1xf64>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f64, %out: f32):
      %883 = arith.truncf %in : f64 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %727 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%717 : tensor<1x32x4x262144xf32>) outs(%21 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.addf %in, %out : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %728 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%727 : tensor<1x32x1x1xf32>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.divf %in, %cst_152 : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %729 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%726 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %730 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%729 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %731 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%728 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %732 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%731 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %733 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%730 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_145 : f64 to f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x32x1x1xf16>
    %734 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%733 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.rsqrt %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %735 = tensor.empty() : tensor<1x32x4x262144xf16>
    %736 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_340, %732 : tensor<1x32x4x262144xf16>, tensor<1x32x1x1xf16>) outs(%735 : tensor<1x32x4x262144xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x4x262144xf16>
    %737 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%736, %734 : tensor<1x32x4x262144xf16>, tensor<1x32x1x1xf16>) outs(%735 : tensor<1x32x4x262144xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x4x262144xf16>
    %collapsed_341 = tensor.collapse_shape %737 [[0], [1, 2], [3]] : tensor<1x32x4x262144xf16> into tensor<1x128x262144xf16>
    %expanded_342 = tensor.expand_shape %collapsed_341 [[0], [1], [2, 3]] : tensor<1x128x262144xf16> into tensor<1x128x512x512xf16>
    %expanded_343 = tensor.expand_shape %cst_24 [[0, 1, 2, 3]] : tensor<128xf16> into tensor<1x128x1x1xf16>
    %expanded_344 = tensor.expand_shape %cst_23 [[0, 1, 2, 3]] : tensor<128xf16> into tensor<1x128x1x1xf16>
    %738 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_342, %expanded_344 : tensor<1x128x512x512xf16>, tensor<1x128x1x1xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x128x512x512xf16>
    %739 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%738, %expanded_343 : tensor<1x128x512x512xf16>, tensor<1x128x1x1xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x128x512x512xf16>
    %740 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%739 : tensor<1x128x512x512xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.negf %in : f16
      %884 = math.exp %883 : f16
      %885 = arith.addf %884, %cst_141 : f16
      %886 = arith.divf %cst_141, %885 : f16
      linalg.yield %886 : f16
    } -> tensor<1x128x512x512xf16>
    %741 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%740, %739 : tensor<1x128x512x512xf16>, tensor<1x128x512x512xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x128x512x512xf16>
    %padded_345 = tensor.pad %741 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x128x512x512xf16> to tensor<1x128x514x514xf16>
    %742 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_21 : tensor<128xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x128x512x512xf16>
    %743 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_345, %cst_22 : tensor<1x128x514x514xf16>, tensor<128x128x3x3xf16>) outs(%742 : tensor<1x128x512x512xf16>) -> tensor<1x128x512x512xf16>
    %744 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_19 : tensor<128xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x128x512x512xf16>
    %745 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%686, %cst_20 : tensor<1x256x512x512xf16>, tensor<128x256x1x1xf16>) outs(%744 : tensor<1x128x512x512xf16>) -> tensor<1x128x512x512xf16>
    %746 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%745, %743 : tensor<1x128x512x512xf16>, tensor<1x128x512x512xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x128x512x512xf16>
    %747 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%746 : tensor<1x128x512x512xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x128x512x512xf16>
    %collapsed_346 = tensor.collapse_shape %747 [[0], [1], [2, 3]] : tensor<1x128x512x512xf16> into tensor<1x128x262144xf16>
    %expanded_347 = tensor.expand_shape %collapsed_346 [[0], [1, 2], [3]] : tensor<1x128x262144xf16> into tensor<1x32x4x262144xf16>
    %748 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_347 : tensor<1x32x4x262144xf16>) outs(%716 : tensor<1x32x4x262144xf32>) {
    ^bb0(%in: f16, %out: f32):
      %883 = arith.extf %in : f16 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x4x262144xf32>
    %749 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%748 : tensor<1x32x4x262144xf32>) outs(%718 : tensor<1x32x4x262144xf64>) {
    ^bb0(%in: f32, %out: f64):
      %883 = arith.extf %in : f32 to f64
      linalg.yield %883 : f64
    } -> tensor<1x32x4x262144xf64>
    %750 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%749 : tensor<1x32x4x262144xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %751 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%750 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_151 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %752 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%749, %751 : tensor<1x32x4x262144xf64>, tensor<1x32x1x1xf64>) outs(%718 : tensor<1x32x4x262144xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.subf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x4x262144xf64>
    %753 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%752, %752 : tensor<1x32x4x262144xf64>, tensor<1x32x4x262144xf64>) outs(%718 : tensor<1x32x4x262144xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.mulf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x4x262144xf64>
    %754 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%753 : tensor<1x32x4x262144xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %755 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%754 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_151 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %756 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%755 : tensor<1x32x1x1xf64>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f64, %out: f32):
      %883 = arith.truncf %in : f64 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %757 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%748 : tensor<1x32x4x262144xf32>) outs(%21 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.addf %in, %out : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %758 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%757 : tensor<1x32x1x1xf32>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.divf %in, %cst_152 : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %759 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%756 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %760 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%759 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %761 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%758 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %762 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%761 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %763 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%760 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_145 : f64 to f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x32x1x1xf16>
    %764 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%763 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.rsqrt %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %765 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_347, %762 : tensor<1x32x4x262144xf16>, tensor<1x32x1x1xf16>) outs(%735 : tensor<1x32x4x262144xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x4x262144xf16>
    %766 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%765, %764 : tensor<1x32x4x262144xf16>, tensor<1x32x1x1xf16>) outs(%735 : tensor<1x32x4x262144xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x4x262144xf16>
    %collapsed_348 = tensor.collapse_shape %766 [[0], [1, 2], [3]] : tensor<1x32x4x262144xf16> into tensor<1x128x262144xf16>
    %expanded_349 = tensor.expand_shape %collapsed_348 [[0], [1], [2, 3]] : tensor<1x128x262144xf16> into tensor<1x128x512x512xf16>
    %expanded_350 = tensor.expand_shape %cst_18 [[0, 1, 2, 3]] : tensor<128xf16> into tensor<1x128x1x1xf16>
    %expanded_351 = tensor.expand_shape %cst_17 [[0, 1, 2, 3]] : tensor<128xf16> into tensor<1x128x1x1xf16>
    %767 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_349, %expanded_351 : tensor<1x128x512x512xf16>, tensor<1x128x1x1xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x128x512x512xf16>
    %768 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%767, %expanded_350 : tensor<1x128x512x512xf16>, tensor<1x128x1x1xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x128x512x512xf16>
    %769 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%768 : tensor<1x128x512x512xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.negf %in : f16
      %884 = math.exp %883 : f16
      %885 = arith.addf %884, %cst_141 : f16
      %886 = arith.divf %cst_141, %885 : f16
      linalg.yield %886 : f16
    } -> tensor<1x128x512x512xf16>
    %770 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%769, %768 : tensor<1x128x512x512xf16>, tensor<1x128x512x512xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x128x512x512xf16>
    %padded_352 = tensor.pad %770 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x128x512x512xf16> to tensor<1x128x514x514xf16>
    %771 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_15 : tensor<128xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x128x512x512xf16>
    %772 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_352, %cst_16 : tensor<1x128x514x514xf16>, tensor<128x128x3x3xf16>) outs(%771 : tensor<1x128x512x512xf16>) -> tensor<1x128x512x512xf16>
    %collapsed_353 = tensor.collapse_shape %772 [[0], [1], [2, 3]] : tensor<1x128x512x512xf16> into tensor<1x128x262144xf16>
    %expanded_354 = tensor.expand_shape %collapsed_353 [[0], [1, 2], [3]] : tensor<1x128x262144xf16> into tensor<1x32x4x262144xf16>
    %773 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_354 : tensor<1x32x4x262144xf16>) outs(%716 : tensor<1x32x4x262144xf32>) {
    ^bb0(%in: f16, %out: f32):
      %883 = arith.extf %in : f16 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x4x262144xf32>
    %774 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%773 : tensor<1x32x4x262144xf32>) outs(%718 : tensor<1x32x4x262144xf64>) {
    ^bb0(%in: f32, %out: f64):
      %883 = arith.extf %in : f32 to f64
      linalg.yield %883 : f64
    } -> tensor<1x32x4x262144xf64>
    %775 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%774 : tensor<1x32x4x262144xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %776 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%775 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_151 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %777 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%774, %776 : tensor<1x32x4x262144xf64>, tensor<1x32x1x1xf64>) outs(%718 : tensor<1x32x4x262144xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.subf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x4x262144xf64>
    %778 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%777, %777 : tensor<1x32x4x262144xf64>, tensor<1x32x4x262144xf64>) outs(%718 : tensor<1x32x4x262144xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.mulf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x4x262144xf64>
    %779 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%778 : tensor<1x32x4x262144xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %780 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%779 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_151 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %781 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%780 : tensor<1x32x1x1xf64>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f64, %out: f32):
      %883 = arith.truncf %in : f64 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %782 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%773 : tensor<1x32x4x262144xf32>) outs(%21 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.addf %in, %out : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %783 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%782 : tensor<1x32x1x1xf32>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.divf %in, %cst_152 : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %784 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%781 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %785 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%784 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %786 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%783 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %787 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%786 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %788 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%785 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_145 : f64 to f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x32x1x1xf16>
    %789 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%788 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.rsqrt %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %790 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_354, %787 : tensor<1x32x4x262144xf16>, tensor<1x32x1x1xf16>) outs(%735 : tensor<1x32x4x262144xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x4x262144xf16>
    %791 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%790, %789 : tensor<1x32x4x262144xf16>, tensor<1x32x1x1xf16>) outs(%735 : tensor<1x32x4x262144xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x4x262144xf16>
    %collapsed_355 = tensor.collapse_shape %791 [[0], [1, 2], [3]] : tensor<1x32x4x262144xf16> into tensor<1x128x262144xf16>
    %expanded_356 = tensor.expand_shape %collapsed_355 [[0], [1], [2, 3]] : tensor<1x128x262144xf16> into tensor<1x128x512x512xf16>
    %expanded_357 = tensor.expand_shape %cst_14 [[0, 1, 2, 3]] : tensor<128xf16> into tensor<1x128x1x1xf16>
    %expanded_358 = tensor.expand_shape %cst_13 [[0, 1, 2, 3]] : tensor<128xf16> into tensor<1x128x1x1xf16>
    %792 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_356, %expanded_358 : tensor<1x128x512x512xf16>, tensor<1x128x1x1xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x128x512x512xf16>
    %793 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%792, %expanded_357 : tensor<1x128x512x512xf16>, tensor<1x128x1x1xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x128x512x512xf16>
    %794 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%793 : tensor<1x128x512x512xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.negf %in : f16
      %884 = math.exp %883 : f16
      %885 = arith.addf %884, %cst_141 : f16
      %886 = arith.divf %cst_141, %885 : f16
      linalg.yield %886 : f16
    } -> tensor<1x128x512x512xf16>
    %795 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%794, %793 : tensor<1x128x512x512xf16>, tensor<1x128x512x512xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x128x512x512xf16>
    %padded_359 = tensor.pad %795 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x128x512x512xf16> to tensor<1x128x514x514xf16>
    %796 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_11 : tensor<128xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x128x512x512xf16>
    %797 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_359, %cst_12 : tensor<1x128x514x514xf16>, tensor<128x128x3x3xf16>) outs(%796 : tensor<1x128x512x512xf16>) -> tensor<1x128x512x512xf16>
    %798 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%747, %797 : tensor<1x128x512x512xf16>, tensor<1x128x512x512xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x128x512x512xf16>
    %799 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%798 : tensor<1x128x512x512xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x128x512x512xf16>
    %collapsed_360 = tensor.collapse_shape %799 [[0], [1], [2, 3]] : tensor<1x128x512x512xf16> into tensor<1x128x262144xf16>
    %expanded_361 = tensor.expand_shape %collapsed_360 [[0], [1, 2], [3]] : tensor<1x128x262144xf16> into tensor<1x32x4x262144xf16>
    %800 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_361 : tensor<1x32x4x262144xf16>) outs(%716 : tensor<1x32x4x262144xf32>) {
    ^bb0(%in: f16, %out: f32):
      %883 = arith.extf %in : f16 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x4x262144xf32>
    %801 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%800 : tensor<1x32x4x262144xf32>) outs(%718 : tensor<1x32x4x262144xf64>) {
    ^bb0(%in: f32, %out: f64):
      %883 = arith.extf %in : f32 to f64
      linalg.yield %883 : f64
    } -> tensor<1x32x4x262144xf64>
    %802 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%801 : tensor<1x32x4x262144xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %803 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%802 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_151 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %804 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%801, %803 : tensor<1x32x4x262144xf64>, tensor<1x32x1x1xf64>) outs(%718 : tensor<1x32x4x262144xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.subf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x4x262144xf64>
    %805 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%804, %804 : tensor<1x32x4x262144xf64>, tensor<1x32x4x262144xf64>) outs(%718 : tensor<1x32x4x262144xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.mulf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x4x262144xf64>
    %806 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%805 : tensor<1x32x4x262144xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %807 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%806 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_151 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %808 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%807 : tensor<1x32x1x1xf64>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f64, %out: f32):
      %883 = arith.truncf %in : f64 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %809 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%800 : tensor<1x32x4x262144xf32>) outs(%21 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.addf %in, %out : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %810 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%809 : tensor<1x32x1x1xf32>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.divf %in, %cst_152 : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %811 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%808 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %812 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%811 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %813 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%810 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %814 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%813 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %815 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%812 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_145 : f64 to f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x32x1x1xf16>
    %816 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%815 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.rsqrt %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %817 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_361, %814 : tensor<1x32x4x262144xf16>, tensor<1x32x1x1xf16>) outs(%735 : tensor<1x32x4x262144xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x4x262144xf16>
    %818 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%817, %816 : tensor<1x32x4x262144xf16>, tensor<1x32x1x1xf16>) outs(%735 : tensor<1x32x4x262144xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x4x262144xf16>
    %collapsed_362 = tensor.collapse_shape %818 [[0], [1, 2], [3]] : tensor<1x32x4x262144xf16> into tensor<1x128x262144xf16>
    %expanded_363 = tensor.expand_shape %collapsed_362 [[0], [1], [2, 3]] : tensor<1x128x262144xf16> into tensor<1x128x512x512xf16>
    %expanded_364 = tensor.expand_shape %cst_10 [[0, 1, 2, 3]] : tensor<128xf16> into tensor<1x128x1x1xf16>
    %expanded_365 = tensor.expand_shape %cst_9 [[0, 1, 2, 3]] : tensor<128xf16> into tensor<1x128x1x1xf16>
    %819 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_363, %expanded_365 : tensor<1x128x512x512xf16>, tensor<1x128x1x1xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x128x512x512xf16>
    %820 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%819, %expanded_364 : tensor<1x128x512x512xf16>, tensor<1x128x1x1xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x128x512x512xf16>
    %821 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%820 : tensor<1x128x512x512xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.negf %in : f16
      %884 = math.exp %883 : f16
      %885 = arith.addf %884, %cst_141 : f16
      %886 = arith.divf %cst_141, %885 : f16
      linalg.yield %886 : f16
    } -> tensor<1x128x512x512xf16>
    %822 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%821, %820 : tensor<1x128x512x512xf16>, tensor<1x128x512x512xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x128x512x512xf16>
    %padded_366 = tensor.pad %822 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x128x512x512xf16> to tensor<1x128x514x514xf16>
    %823 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_7 : tensor<128xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x128x512x512xf16>
    %824 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_366, %cst_8 : tensor<1x128x514x514xf16>, tensor<128x128x3x3xf16>) outs(%823 : tensor<1x128x512x512xf16>) -> tensor<1x128x512x512xf16>
    %collapsed_367 = tensor.collapse_shape %824 [[0], [1], [2, 3]] : tensor<1x128x512x512xf16> into tensor<1x128x262144xf16>
    %expanded_368 = tensor.expand_shape %collapsed_367 [[0], [1, 2], [3]] : tensor<1x128x262144xf16> into tensor<1x32x4x262144xf16>
    %825 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_368 : tensor<1x32x4x262144xf16>) outs(%716 : tensor<1x32x4x262144xf32>) {
    ^bb0(%in: f16, %out: f32):
      %883 = arith.extf %in : f16 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x4x262144xf32>
    %826 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%825 : tensor<1x32x4x262144xf32>) outs(%718 : tensor<1x32x4x262144xf64>) {
    ^bb0(%in: f32, %out: f64):
      %883 = arith.extf %in : f32 to f64
      linalg.yield %883 : f64
    } -> tensor<1x32x4x262144xf64>
    %827 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%826 : tensor<1x32x4x262144xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %828 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%827 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_151 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %829 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%826, %828 : tensor<1x32x4x262144xf64>, tensor<1x32x1x1xf64>) outs(%718 : tensor<1x32x4x262144xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.subf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x4x262144xf64>
    %830 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%829, %829 : tensor<1x32x4x262144xf64>, tensor<1x32x4x262144xf64>) outs(%718 : tensor<1x32x4x262144xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.mulf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x4x262144xf64>
    %831 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%830 : tensor<1x32x4x262144xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %832 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%831 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_151 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %833 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%832 : tensor<1x32x1x1xf64>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f64, %out: f32):
      %883 = arith.truncf %in : f64 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %834 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%825 : tensor<1x32x4x262144xf32>) outs(%21 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.addf %in, %out : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %835 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%834 : tensor<1x32x1x1xf32>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.divf %in, %cst_152 : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %836 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%833 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %837 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%836 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %838 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%835 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %839 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%838 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %840 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%837 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_145 : f64 to f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x32x1x1xf16>
    %841 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%840 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.rsqrt %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %842 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_368, %839 : tensor<1x32x4x262144xf16>, tensor<1x32x1x1xf16>) outs(%735 : tensor<1x32x4x262144xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x4x262144xf16>
    %843 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%842, %841 : tensor<1x32x4x262144xf16>, tensor<1x32x1x1xf16>) outs(%735 : tensor<1x32x4x262144xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x4x262144xf16>
    %collapsed_369 = tensor.collapse_shape %843 [[0], [1, 2], [3]] : tensor<1x32x4x262144xf16> into tensor<1x128x262144xf16>
    %expanded_370 = tensor.expand_shape %collapsed_369 [[0], [1], [2, 3]] : tensor<1x128x262144xf16> into tensor<1x128x512x512xf16>
    %expanded_371 = tensor.expand_shape %cst_6 [[0, 1, 2, 3]] : tensor<128xf16> into tensor<1x128x1x1xf16>
    %expanded_372 = tensor.expand_shape %cst_5 [[0, 1, 2, 3]] : tensor<128xf16> into tensor<1x128x1x1xf16>
    %844 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_370, %expanded_372 : tensor<1x128x512x512xf16>, tensor<1x128x1x1xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x128x512x512xf16>
    %845 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%844, %expanded_371 : tensor<1x128x512x512xf16>, tensor<1x128x1x1xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x128x512x512xf16>
    %846 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%845 : tensor<1x128x512x512xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.negf %in : f16
      %884 = math.exp %883 : f16
      %885 = arith.addf %884, %cst_141 : f16
      %886 = arith.divf %cst_141, %885 : f16
      linalg.yield %886 : f16
    } -> tensor<1x128x512x512xf16>
    %847 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%846, %845 : tensor<1x128x512x512xf16>, tensor<1x128x512x512xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x128x512x512xf16>
    %padded_373 = tensor.pad %847 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x128x512x512xf16> to tensor<1x128x514x514xf16>
    %848 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst_3 : tensor<128xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x128x512x512xf16>
    %849 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_373, %cst_4 : tensor<1x128x514x514xf16>, tensor<128x128x3x3xf16>) outs(%848 : tensor<1x128x512x512xf16>) -> tensor<1x128x512x512xf16>
    %850 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%799, %849 : tensor<1x128x512x512xf16>, tensor<1x128x512x512xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x128x512x512xf16>
    %851 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%850 : tensor<1x128x512x512xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x128x512x512xf16>
    %collapsed_374 = tensor.collapse_shape %851 [[0], [1], [2, 3]] : tensor<1x128x512x512xf16> into tensor<1x128x262144xf16>
    %expanded_375 = tensor.expand_shape %collapsed_374 [[0], [1, 2], [3]] : tensor<1x128x262144xf16> into tensor<1x32x4x262144xf16>
    %852 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_375 : tensor<1x32x4x262144xf16>) outs(%716 : tensor<1x32x4x262144xf32>) {
    ^bb0(%in: f16, %out: f32):
      %883 = arith.extf %in : f16 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x4x262144xf32>
    %853 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%852 : tensor<1x32x4x262144xf32>) outs(%718 : tensor<1x32x4x262144xf64>) {
    ^bb0(%in: f32, %out: f64):
      %883 = arith.extf %in : f32 to f64
      linalg.yield %883 : f64
    } -> tensor<1x32x4x262144xf64>
    %854 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%853 : tensor<1x32x4x262144xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %855 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%854 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_151 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %856 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%853, %855 : tensor<1x32x4x262144xf64>, tensor<1x32x1x1xf64>) outs(%718 : tensor<1x32x4x262144xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.subf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x4x262144xf64>
    %857 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%856, %856 : tensor<1x32x4x262144xf64>, tensor<1x32x4x262144xf64>) outs(%718 : tensor<1x32x4x262144xf64>) {
    ^bb0(%in: f64, %in_381: f64, %out: f64):
      %883 = arith.mulf %in, %in_381 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x4x262144xf64>
    %858 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%857 : tensor<1x32x4x262144xf64>) outs(%12 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.addf %in, %out : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %859 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%858 : tensor<1x32x1x1xf64>) outs(%11 : tensor<1x32x1x1xf64>) {
    ^bb0(%in: f64, %out: f64):
      %883 = arith.divf %in, %cst_151 : f64
      linalg.yield %883 : f64
    } -> tensor<1x32x1x1xf64>
    %860 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%859 : tensor<1x32x1x1xf64>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f64, %out: f32):
      %883 = arith.truncf %in : f64 to f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %861 = linalg.generic {indexing_maps = [#map1, #map3], iterator_types = ["parallel", "parallel", "reduction", "reduction"]} ins(%852 : tensor<1x32x4x262144xf32>) outs(%21 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.addf %in, %out : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %862 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%861 : tensor<1x32x1x1xf32>) outs(%19 : tensor<1x32x1x1xf32>) {
    ^bb0(%in: f32, %out: f32):
      %883 = arith.divf %in, %cst_152 : f32
      linalg.yield %883 : f32
    } -> tensor<1x32x1x1xf32>
    %863 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%860 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %864 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%863 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %865 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%862 : tensor<1x32x1x1xf32>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f32, %out: f16):
      %883 = arith.truncf %in : f32 to f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %866 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%865 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x32x1x1xf16>
    %867 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%864 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.truncf %cst_145 : f64 to f16
      %884 = arith.addf %in, %883 : f16
      linalg.yield %884 : f16
    } -> tensor<1x32x1x1xf16>
    %868 = linalg.generic {indexing_maps = [#map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%867 : tensor<1x32x1x1xf16>) outs(%24 : tensor<1x32x1x1xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.rsqrt %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x1x1xf16>
    %869 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_375, %866 : tensor<1x32x4x262144xf16>, tensor<1x32x1x1xf16>) outs(%735 : tensor<1x32x4x262144xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.subf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x4x262144xf16>
    %870 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%869, %868 : tensor<1x32x4x262144xf16>, tensor<1x32x1x1xf16>) outs(%735 : tensor<1x32x4x262144xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x32x4x262144xf16>
    %collapsed_376 = tensor.collapse_shape %870 [[0], [1, 2], [3]] : tensor<1x32x4x262144xf16> into tensor<1x128x262144xf16>
    %expanded_377 = tensor.expand_shape %collapsed_376 [[0], [1], [2, 3]] : tensor<1x128x262144xf16> into tensor<1x128x512x512xf16>
    %expanded_378 = tensor.expand_shape %cst_2 [[0, 1, 2, 3]] : tensor<128xf16> into tensor<1x128x1x1xf16>
    %expanded_379 = tensor.expand_shape %cst_1 [[0, 1, 2, 3]] : tensor<128xf16> into tensor<1x128x1x1xf16>
    %871 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%expanded_377, %expanded_379 : tensor<1x128x512x512xf16>, tensor<1x128x1x1xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x128x512x512xf16>
    %872 = linalg.generic {indexing_maps = [#map, #map4, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%871, %expanded_378 : tensor<1x128x512x512xf16>, tensor<1x128x1x1xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.addf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x128x512x512xf16>
    %873 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%872 : tensor<1x128x512x512xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.negf %in : f16
      %884 = math.exp %883 : f16
      %885 = arith.addf %884, %cst_141 : f16
      %886 = arith.divf %cst_141, %885 : f16
      linalg.yield %886 : f16
    } -> tensor<1x128x512x512xf16>
    %874 = linalg.generic {indexing_maps = [#map, #map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%873, %872 : tensor<1x128x512x512xf16>, tensor<1x128x512x512xf16>) outs(%713 : tensor<1x128x512x512xf16>) {
    ^bb0(%in: f16, %in_381: f16, %out: f16):
      %883 = arith.mulf %in, %in_381 : f16
      linalg.yield %883 : f16
    } -> tensor<1x128x512x512xf16>
    %padded_380 = tensor.pad %874 low[0, 0, 1, 1] high[0, 0, 1, 1] {
    ^bb0(%arg1: index, %arg2: index, %arg3: index, %arg4: index):
      tensor.yield %cst_139 : f16
    } : tensor<1x128x512x512xf16> to tensor<1x128x514x514xf16>
    %875 = tensor.empty() : tensor<1x3x512x512xf16>
    %876 = linalg.generic {indexing_maps = [#map2, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%cst : tensor<3xf16>) outs(%875 : tensor<1x3x512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      linalg.yield %in : f16
    } -> tensor<1x3x512x512xf16>
    %877 = linalg.conv_2d_nchw_fchw {dilations = dense<1> : vector<2xi64>, strides = dense<1> : vector<2xi64>} ins(%padded_380, %cst_0 : tensor<1x128x514x514xf16>, tensor<3x128x3x3xf16>) outs(%876 : tensor<1x3x512x512xf16>) -> tensor<1x3x512x512xf16>
    %878 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%877 : tensor<1x3x512x512xf16>) outs(%875 : tensor<1x3x512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.divf %in, %cst_157 : f16
      linalg.yield %883 : f16
    } -> tensor<1x3x512x512xf16>
    %879 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%878 : tensor<1x3x512x512xf16>) outs(%875 : tensor<1x3x512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.addf %in, %cst_158 : f16
      linalg.yield %883 : f16
    } -> tensor<1x3x512x512xf16>
    %880 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%879 : tensor<1x3x512x512xf16>) outs(%875 : tensor<1x3x512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.cmpf ult, %in, %cst_139 : f16
      %884 = arith.select %883, %cst_139, %in : f16
      %885 = arith.cmpf ugt, %884, %cst_141 : f16
      %886 = arith.select %885, %cst_141, %884 : f16
      linalg.yield %886 : f16
    } -> tensor<1x3x512x512xf16>
    %881 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%880 : tensor<1x3x512x512xf16>) outs(%875 : tensor<1x3x512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = arith.mulf %in, %cst_159 : f16
      linalg.yield %883 : f16
    } -> tensor<1x3x512x512xf16>
    %882 = linalg.generic {indexing_maps = [#map, #map1], iterator_types = ["parallel", "parallel", "parallel", "parallel"]} ins(%881 : tensor<1x3x512x512xf16>) outs(%875 : tensor<1x3x512x512xf16>) {
    ^bb0(%in: f16, %out: f16):
      %883 = math.roundeven %in : f16
      linalg.yield %883 : f16
    } -> tensor<1x3x512x512xf16>
    return %882 : tensor<1x3x512x512xf16>
  }
}