High Level JAX export

mnist_train.mlir

module  {
  iree_input.global private mutable @init_params$0 : tensor<784x1024xf32>
  iree_input.global private mutable @init_params$1 : tensor<1024xf32>
  iree_input.global private mutable @init_params$2 : tensor<1024x1024xf32>
  iree_input.global private mutable @init_params$3 : tensor<1024xf32>
  iree_input.global private mutable @init_params$4 : tensor<1024x10xf32>
  iree_input.global private mutable @init_params$5 : tensor<10xf32>
  iree_input.global private mutable @opt_state$1 : tensor<784x1024xf32>
  iree_input.global private mutable @opt_state$3 : tensor<1024xf32>
  iree_input.global private mutable @opt_state$5 : tensor<1024x1024xf32>
  iree_input.global private mutable @opt_state$7 : tensor<1024xf32>
  iree_input.global private mutable @opt_state$9 : tensor<1024x10xf32>
  iree_input.global private mutable @opt_state$11 : tensor<10xf32>
  func @get_params() -> (tensor<784x1024xf32>, tensor<1024xf32>, tensor<1024x1024xf32>, tensor<1024xf32>, tensor<1024x10xf32>, tensor<10xf32>) {
    %0 = iree_input.global.load @init_params$0 : tensor<784x1024xf32>
    %1 = iree_input.global.load @init_params$1 : tensor<1024xf32>
    %2 = iree_input.global.load @init_params$2 : tensor<1024x1024xf32>
    %3 = iree_input.global.load @init_params$3 : tensor<1024xf32>
    %4 = iree_input.global.load @init_params$4 : tensor<1024x10xf32>
    %5 = iree_input.global.load @init_params$5 : tensor<10xf32>
    return %0, %1, %2, %3, %4, %5 : tensor<784x1024xf32>, tensor<1024xf32>, tensor<1024x1024xf32>, tensor<1024xf32>, tensor<1024x10xf32>, tensor<10xf32>
  }
  func @get_opt_state() -> (tensor<784x1024xf32>, tensor<784x1024xf32>, tensor<1024xf32>, tensor<1024xf32>, tensor<1024x1024xf32>, tensor<1024x1024xf32>, tensor<1024xf32>, tensor<1024xf32>, tensor<1024x10xf32>, tensor<1024x10xf32>, tensor<10xf32>, tensor<10xf32>) {
    %0 = iree_input.global.load @init_params$0 : tensor<784x1024xf32>
    %1 = iree_input.global.load @opt_state$1 : tensor<784x1024xf32>
    %2 = iree_input.global.load @init_params$1 : tensor<1024xf32>
    %3 = iree_input.global.load @opt_state$3 : tensor<1024xf32>
    %4 = iree_input.global.load @init_params$2 : tensor<1024x1024xf32>
    %5 = iree_input.global.load @opt_state$5 : tensor<1024x1024xf32>
    %6 = iree_input.global.load @init_params$3 : tensor<1024xf32>
    %7 = iree_input.global.load @opt_state$7 : tensor<1024xf32>
    %8 = iree_input.global.load @init_params$4 : tensor<1024x10xf32>
    %9 = iree_input.global.load @opt_state$9 : tensor<1024x10xf32>
    %10 = iree_input.global.load @init_params$5 : tensor<10xf32>
    %11 = iree_input.global.load @opt_state$11 : tensor<10xf32>
    return %0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11 : tensor<784x1024xf32>, tensor<784x1024xf32>, tensor<1024xf32>, tensor<1024xf32>, tensor<1024x1024xf32>, tensor<1024x1024xf32>, tensor<1024xf32>, tensor<1024xf32>, tensor<1024x10xf32>, tensor<1024x10xf32>, tensor<10xf32>, tensor<10xf32>
  }
  func @set_opt_state(%arg0: tensor<784x1024xf32>, %arg1: tensor<784x1024xf32>, %arg2: tensor<1024xf32>, %arg3: tensor<1024xf32>, %arg4: tensor<1024x1024xf32>, %arg5: tensor<1024x1024xf32>, %arg6: tensor<1024xf32>, %arg7: tensor<1024xf32>, %arg8: tensor<1024x10xf32>, %arg9: tensor<1024x10xf32>, %arg10: tensor<10xf32>, %arg11: tensor<10xf32>) {
    iree_input.global.store %arg0, @init_params$0 : tensor<784x1024xf32>
    iree_input.global.store %arg1, @opt_state$1 : tensor<784x1024xf32>
    iree_input.global.store %arg2, @init_params$1 : tensor<1024xf32>
    iree_input.global.store %arg3, @opt_state$3 : tensor<1024xf32>
    iree_input.global.store %arg4, @init_params$2 : tensor<1024x1024xf32>
    iree_input.global.store %arg5, @opt_state$5 : tensor<1024x1024xf32>
    iree_input.global.store %arg6, @init_params$3 : tensor<1024xf32>
    iree_input.global.store %arg7, @opt_state$7 : tensor<1024xf32>
    iree_input.global.store %arg8, @init_params$4 : tensor<1024x10xf32>
    iree_input.global.store %arg9, @opt_state$9 : tensor<1024x10xf32>
    iree_input.global.store %arg10, @init_params$5 : tensor<10xf32>
    iree_input.global.store %arg11, @opt_state$11 : tensor<10xf32>
    return
  }
  func @initialize(%arg0: tensor<2xui32>) {
    %0:12 = call @main(%arg0) : (tensor<2xui32>) -> (tensor<784x1024xf32>, tensor<784x1024xf32>, tensor<1024xf32>, tensor<1024xf32>, tensor<1024x1024xf32>, tensor<1024x1024xf32>, tensor<1024xf32>, tensor<1024xf32>, tensor<1024x10xf32>, tensor<1024x10xf32>, tensor<10xf32>, tensor<10xf32>)
    iree_input.global.store %0#0, @init_params$0 : tensor<784x1024xf32>
    iree_input.global.store %0#1, @opt_state$1 : tensor<784x1024xf32>
    iree_input.global.store %0#2, @init_params$1 : tensor<1024xf32>
    iree_input.global.store %0#3, @opt_state$3 : tensor<1024xf32>
    iree_input.global.store %0#4, @init_params$2 : tensor<1024x1024xf32>
    iree_input.global.store %0#5, @opt_state$5 : tensor<1024x1024xf32>
    iree_input.global.store %0#6, @init_params$3 : tensor<1024xf32>
    iree_input.global.store %0#7, @opt_state$7 : tensor<1024xf32>
    iree_input.global.store %0#8, @init_params$4 : tensor<1024x10xf32>
    iree_input.global.store %0#9, @opt_state$9 : tensor<1024x10xf32>
    iree_input.global.store %0#10, @init_params$5 : tensor<10xf32>
    iree_input.global.store %0#11, @opt_state$11 : tensor<10xf32>
    return
  }
  func private @main(%arg0: tensor<2xui32>) -> (tensor<784x1024xf32>, tensor<784x1024xf32>, tensor<1024xf32>, tensor<1024xf32>, tensor<1024x1024xf32>, tensor<1024x1024xf32>, tensor<1024xf32>, tensor<1024xf32>, tensor<1024x10xf32>, tensor<1024x10xf32>, tensor<10xf32>, tensor<10xf32>) {
    %0 = mhlo.constant dense<0.000000e+00> : tensor<10xf32>
    %1 = mhlo.constant dense<0.000000e+00> : tensor<1024x10xf32>
    %2 = mhlo.constant dense<0.000000e+00> : tensor<1024xf32>
    %3 = mhlo.constant dense<0.000000e+00> : tensor<1024x1024xf32>
    %4 = mhlo.constant dense<0.000000e+00> : tensor<784x1024xf32>
    %5 = mhlo.constant dense<5> : tensor<2xui32>
    %6 = mhlo.constant dense<26> : tensor<2xui32>
    %7 = mhlo.constant dense<6> : tensor<2xui32>
    %8 = mhlo.constant dense<17> : tensor<2xui32>
    %9 = mhlo.constant dense<15> : tensor<2xui32>
    %10 = mhlo.constant dense<19> : tensor<2xui32>
    %11 = mhlo.constant dense<13> : tensor<2xui32>
    %12 = mhlo.constant dense<4> : tensor<2xui32>
    %13 = mhlo.constant dense<8> : tensor<2xui32>
    %14 = mhlo.constant dense<24> : tensor<2xui32>
    %15 = mhlo.constant dense<16> : tensor<2xui32>
    %16 = mhlo.constant dense<3> : tensor<2xui32>
    %17 = mhlo.constant dense<29> : tensor<2xui32>
    %18 = mhlo.constant dense<2> : tensor<2xui32>
    %19 = mhlo.constant dense<1> : tensor<2xui32>
    %20 = mhlo.constant dense<466688986> : tensor<ui32>
    %21 = mhlo.constant dense<0.00999999977> : tensor<10xf32>
    %22 = mhlo.constant dense<1.41421354> : tensor<10xf32>
    %23 = mhlo.constant dense<1.41421354> : tensor<f32>
    %24 = mhlo.constant dense<-0.99999994> : tensor<10xf32>
    %25 = mhlo.constant dense<2.000000e+00> : tensor<10xf32>
    %26 = mhlo.constant dense<1.000000e+00> : tensor<10xf32>
    %27 = mhlo.constant dense<1065353216> : tensor<10xui32>
    %28 = mhlo.constant dense<9> : tensor<10xui32>
    %29 = mhlo.constant dense<5> : tensor<5xui32>
    %30 = mhlo.constant dense<26> : tensor<5xui32>
    %31 = mhlo.constant dense<6> : tensor<5xui32>
    %32 = mhlo.constant dense<17> : tensor<5xui32>
    %33 = mhlo.constant dense<15> : tensor<5xui32>
    %34 = mhlo.constant dense<19> : tensor<5xui32>
    %35 = mhlo.constant dense<13> : tensor<5xui32>
    %36 = mhlo.constant dense<4> : tensor<5xui32>
    %37 = mhlo.constant dense<8> : tensor<5xui32>
    %38 = mhlo.constant dense<24> : tensor<5xui32>
    %39 = mhlo.constant dense<16> : tensor<5xui32>
    %40 = mhlo.constant dense<3> : tensor<5xui32>
    %41 = mhlo.constant dense<29> : tensor<5xui32>
    %42 = mhlo.constant dense<2> : tensor<5xui32>
    %43 = mhlo.constant dense<1> : tensor<5xui32>
    %44 = mhlo.constant dense<0.0499984846> : tensor<1024x10xf32>
    %45 = mhlo.constant dense<0xFF800000> : tensor<f32>
    %46 = mhlo.constant dense<2.000000e+00> : tensor<f32>
    %47 = mhlo.constant dense<0x7F800000> : tensor<f32>
    %48 = mhlo.constant dense<-2.000000e+00> : tensor<f32>
    %49 = mhlo.constant dense<1.41421354> : tensor<1024x10xf32>
    %50 = mhlo.constant dense<1.000000e+00> : tensor<1024x10xf32>
    %51 = mhlo.constant dense<1065353216> : tensor<1024x10xui32>
    %52 = mhlo.constant dense<9> : tensor<1024x10xui32>
    %53 = mhlo.constant dense<5> : tensor<5120xui32>
    %54 = mhlo.constant dense<26> : tensor<5120xui32>
    %55 = mhlo.constant dense<6> : tensor<5120xui32>
    %56 = mhlo.constant dense<17> : tensor<5120xui32>
    %57 = mhlo.constant dense<15> : tensor<5120xui32>
    %58 = mhlo.constant dense<19> : tensor<5120xui32>
    %59 = mhlo.constant dense<13> : tensor<5120xui32>
    %60 = mhlo.constant dense<4> : tensor<5120xui32>
    %61 = mhlo.constant dense<8> : tensor<5120xui32>
    %62 = mhlo.constant dense<24> : tensor<5120xui32>
    %63 = mhlo.constant dense<16> : tensor<5120xui32>
    %64 = mhlo.constant dense<3> : tensor<5120xui32>
    %65 = mhlo.constant dense<29> : tensor<5120xui32>
    %66 = mhlo.constant dense<2> : tensor<5120xui32>
    %67 = mhlo.constant dense<1> : tensor<5120xui32>
    %68 = mhlo.constant dense<-1.41421354> : tensor<f32>
    %69 = mhlo.constant dense<0.00999999977> : tensor<1024xf32>
    %70 = mhlo.constant dense<1.41421354> : tensor<1024xf32>
    %71 = mhlo.constant dense<-0.99999994> : tensor<1024xf32>
    %72 = mhlo.constant dense<2.000000e+00> : tensor<1024xf32>
    %73 = mhlo.constant dense<1.000000e+00> : tensor<1024xf32>
    %74 = mhlo.constant dense<1065353216> : tensor<1024xui32>
    %75 = mhlo.constant dense<9> : tensor<1024xui32>
    %76 = mhlo.constant dense<5> : tensor<512xui32>
    %77 = mhlo.constant dense<26> : tensor<512xui32>
    %78 = mhlo.constant dense<6> : tensor<512xui32>
    %79 = mhlo.constant dense<17> : tensor<512xui32>
    %80 = mhlo.constant dense<15> : tensor<512xui32>
    %81 = mhlo.constant dense<19> : tensor<512xui32>
    %82 = mhlo.constant dense<13> : tensor<512xui32>
    %83 = mhlo.constant dense<4> : tensor<512xui32>
    %84 = mhlo.constant dense<8> : tensor<512xui32>
    %85 = mhlo.constant dense<24> : tensor<512xui32>
    %86 = mhlo.constant dense<16> : tensor<512xui32>
    %87 = mhlo.constant dense<3> : tensor<512xui32>
    %88 = mhlo.constant dense<29> : tensor<512xui32>
    %89 = mhlo.constant dense<2> : tensor<512xui32>
    %90 = mhlo.constant dense<1> : tensor<512xui32>
    %91 = mhlo.constant dense<0.0355264768> : tensor<1024x1024xf32>
    %92 = mhlo.constant dense<1.41421354> : tensor<1024x1024xf32>
    %93 = mhlo.constant dense<1.000000e+00> : tensor<1024x1024xf32>
    %94 = mhlo.constant dense<1065353216> : tensor<1024x1024xui32>
    %95 = mhlo.constant dense<9> : tensor<1024x1024xui32>
    %96 = mhlo.constant dense<5> : tensor<524288xui32>
    %97 = mhlo.constant dense<26> : tensor<524288xui32>
    %98 = mhlo.constant dense<6> : tensor<524288xui32>
    %99 = mhlo.constant dense<17> : tensor<524288xui32>
    %100 = mhlo.constant dense<15> : tensor<524288xui32>
    %101 = mhlo.constant dense<19> : tensor<524288xui32>
    %102 = mhlo.constant dense<13> : tensor<524288xui32>
    %103 = mhlo.constant dense<4> : tensor<524288xui32>
    %104 = mhlo.constant dense<8> : tensor<524288xui32>
    %105 = mhlo.constant dense<24> : tensor<524288xui32>
    %106 = mhlo.constant dense<16> : tensor<524288xui32>
    %107 = mhlo.constant dense<3> : tensor<524288xui32>
    %108 = mhlo.constant dense<29> : tensor<524288xui32>
    %109 = mhlo.constant dense<2> : tensor<524288xui32>
    %110 = mhlo.constant dense<1> : tensor<524288xui32>
    %111 = mhlo.constant dense<0.0378109738> : tensor<784x1024xf32>
    %112 = mhlo.constant dense<1.41421354> : tensor<784x1024xf32>
    %113 = mhlo.constant dense<1.000000e+00> : tensor<784x1024xf32>
    %114 = mhlo.constant dense<1065353216> : tensor<784x1024xui32>
    %115 = mhlo.constant dense<9> : tensor<784x1024xui32>
    %116 = mhlo.constant dense<5> : tensor<401408xui32>
    %117 = mhlo.constant dense<26> : tensor<401408xui32>
    %118 = mhlo.constant dense<6> : tensor<401408xui32>
    %119 = mhlo.constant dense<17> : tensor<401408xui32>
    %120 = mhlo.constant dense<15> : tensor<401408xui32>
    %121 = mhlo.constant dense<19> : tensor<401408xui32>
    %122 = mhlo.constant dense<13> : tensor<401408xui32>
    %123 = mhlo.constant dense<4> : tensor<401408xui32>
    %124 = mhlo.constant dense<8> : tensor<401408xui32>
    %125 = mhlo.constant dense<24> : tensor<401408xui32>
    %126 = mhlo.constant dense<16> : tensor<401408xui32>
    %127 = mhlo.constant dense<3> : tensor<401408xui32>
    %128 = mhlo.constant dense<29> : tensor<401408xui32>
    %129 = mhlo.constant dense<2> : tensor<401408xui32>
    %130 = mhlo.constant dense<1> : tensor<401408xui32>
    %131 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<4xui32>
    %132 = "mhlo.slice"(%arg0) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xui32>) -> tensor<1xui32>
    %133 = "mhlo.reshape"(%132) : (tensor<1xui32>) -> tensor<ui32>
    %134 = "mhlo.slice"(%arg0) {limit_indices = dense<2> : tensor<1xi64>, start_indices = dense<1> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xui32>) -> tensor<1xui32>
    %135 = "mhlo.reshape"(%134) : (tensor<1xui32>) -> tensor<ui32>
    %136 = "mhlo.slice"(%131) {limit_indices = dense<2> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<4xui32>) -> tensor<2xui32>
    %137 = "mhlo.slice"(%131) {limit_indices = dense<4> : tensor<1xi64>, start_indices = dense<2> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<4xui32>) -> tensor<2xui32>
    %138 = mhlo.xor %133, %135 : tensor<ui32>
    %139 = mhlo.xor %138, %20 : tensor<ui32>
    %140 = "mhlo.broadcast_in_dim"(%133) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %141 = mhlo.add %136, %140 : tensor<2xui32>
    %142 = "mhlo.broadcast_in_dim"(%135) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %143 = mhlo.add %137, %142 : tensor<2xui32>
    %144 = mhlo.add %141, %143 : tensor<2xui32>
    %145 = mhlo.shift_left %143, %11 : tensor<2xui32>
    %146 = mhlo.shift_right_logical %143, %10 : tensor<2xui32>
    %147 = mhlo.or %145, %146 : tensor<2xui32>
    %148 = mhlo.xor %144, %147 : tensor<2xui32>
    %149 = mhlo.add %144, %148 : tensor<2xui32>
    %150 = mhlo.shift_left %148, %9 : tensor<2xui32>
    %151 = mhlo.shift_right_logical %148, %8 : tensor<2xui32>
    %152 = mhlo.or %150, %151 : tensor<2xui32>
    %153 = mhlo.xor %149, %152 : tensor<2xui32>
    %154 = mhlo.add %149, %153 : tensor<2xui32>
    %155 = mhlo.shift_left %153, %6 : tensor<2xui32>
    %156 = mhlo.shift_right_logical %153, %7 : tensor<2xui32>
    %157 = mhlo.or %155, %156 : tensor<2xui32>
    %158 = mhlo.xor %154, %157 : tensor<2xui32>
    %159 = mhlo.add %154, %158 : tensor<2xui32>
    %160 = mhlo.shift_left %158, %7 : tensor<2xui32>
    %161 = mhlo.shift_right_logical %158, %6 : tensor<2xui32>
    %162 = mhlo.or %160, %161 : tensor<2xui32>
    %163 = mhlo.xor %159, %162 : tensor<2xui32>
    %164 = "mhlo.broadcast_in_dim"(%135) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %165 = mhlo.add %159, %164 : tensor<2xui32>
    %166 = "mhlo.broadcast_in_dim"(%139) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %167 = mhlo.add %163, %166 : tensor<2xui32>
    %168 = mhlo.add %167, %19 : tensor<2xui32>
    %169 = mhlo.add %165, %168 : tensor<2xui32>
    %170 = mhlo.shift_left %168, %8 : tensor<2xui32>
    %171 = mhlo.shift_right_logical %168, %9 : tensor<2xui32>
    %172 = mhlo.or %170, %171 : tensor<2xui32>
    %173 = mhlo.xor %169, %172 : tensor<2xui32>
    %174 = mhlo.add %169, %173 : tensor<2xui32>
    %175 = mhlo.shift_left %173, %17 : tensor<2xui32>
    %176 = mhlo.shift_right_logical %173, %16 : tensor<2xui32>
    %177 = mhlo.or %175, %176 : tensor<2xui32>
    %178 = mhlo.xor %174, %177 : tensor<2xui32>
    %179 = mhlo.add %174, %178 : tensor<2xui32>
    %180 = mhlo.shift_left %178, %15 : tensor<2xui32>
    %181 = mhlo.shift_right_logical %178, %15 : tensor<2xui32>
    %182 = mhlo.or %180, %181 : tensor<2xui32>
    %183 = mhlo.xor %179, %182 : tensor<2xui32>
    %184 = mhlo.add %179, %183 : tensor<2xui32>
    %185 = mhlo.shift_left %183, %14 : tensor<2xui32>
    %186 = mhlo.shift_right_logical %183, %13 : tensor<2xui32>
    %187 = mhlo.or %185, %186 : tensor<2xui32>
    %188 = mhlo.xor %184, %187 : tensor<2xui32>
    %189 = "mhlo.broadcast_in_dim"(%139) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %190 = mhlo.add %184, %189 : tensor<2xui32>
    %191 = "mhlo.broadcast_in_dim"(%133) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %192 = mhlo.add %188, %191 : tensor<2xui32>
    %193 = mhlo.add %192, %18 : tensor<2xui32>
    %194 = mhlo.add %190, %193 : tensor<2xui32>
    %195 = mhlo.shift_left %193, %11 : tensor<2xui32>
    %196 = mhlo.shift_right_logical %193, %10 : tensor<2xui32>
    %197 = mhlo.or %195, %196 : tensor<2xui32>
    %198 = mhlo.xor %194, %197 : tensor<2xui32>
    %199 = mhlo.add %194, %198 : tensor<2xui32>
    %200 = mhlo.shift_left %198, %9 : tensor<2xui32>
    %201 = mhlo.shift_right_logical %198, %8 : tensor<2xui32>
    %202 = mhlo.or %200, %201 : tensor<2xui32>
    %203 = mhlo.xor %199, %202 : tensor<2xui32>
    %204 = mhlo.add %199, %203 : tensor<2xui32>
    %205 = mhlo.shift_left %203, %6 : tensor<2xui32>
    %206 = mhlo.shift_right_logical %203, %7 : tensor<2xui32>
    %207 = mhlo.or %205, %206 : tensor<2xui32>
    %208 = mhlo.xor %204, %207 : tensor<2xui32>
    %209 = mhlo.add %204, %208 : tensor<2xui32>
    %210 = mhlo.shift_left %208, %7 : tensor<2xui32>
    %211 = mhlo.shift_right_logical %208, %6 : tensor<2xui32>
    %212 = mhlo.or %210, %211 : tensor<2xui32>
    %213 = mhlo.xor %209, %212 : tensor<2xui32>
    %214 = "mhlo.broadcast_in_dim"(%133) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %215 = mhlo.add %209, %214 : tensor<2xui32>
    %216 = "mhlo.broadcast_in_dim"(%135) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %217 = mhlo.add %213, %216 : tensor<2xui32>
    %218 = mhlo.add %217, %16 : tensor<2xui32>
    %219 = mhlo.add %215, %218 : tensor<2xui32>
    %220 = mhlo.shift_left %218, %8 : tensor<2xui32>
    %221 = mhlo.shift_right_logical %218, %9 : tensor<2xui32>
    %222 = mhlo.or %220, %221 : tensor<2xui32>
    %223 = mhlo.xor %219, %222 : tensor<2xui32>
    %224 = mhlo.add %219, %223 : tensor<2xui32>
    %225 = mhlo.shift_left %223, %17 : tensor<2xui32>
    %226 = mhlo.shift_right_logical %223, %16 : tensor<2xui32>
    %227 = mhlo.or %225, %226 : tensor<2xui32>
    %228 = mhlo.xor %224, %227 : tensor<2xui32>
    %229 = mhlo.add %224, %228 : tensor<2xui32>
    %230 = mhlo.shift_left %228, %15 : tensor<2xui32>
    %231 = mhlo.shift_right_logical %228, %15 : tensor<2xui32>
    %232 = mhlo.or %230, %231 : tensor<2xui32>
    %233 = mhlo.xor %229, %232 : tensor<2xui32>
    %234 = mhlo.add %229, %233 : tensor<2xui32>
    %235 = mhlo.shift_left %233, %14 : tensor<2xui32>
    %236 = mhlo.shift_right_logical %233, %13 : tensor<2xui32>
    %237 = mhlo.or %235, %236 : tensor<2xui32>
    %238 = mhlo.xor %234, %237 : tensor<2xui32>
    %239 = "mhlo.broadcast_in_dim"(%135) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %240 = mhlo.add %234, %239 : tensor<2xui32>
    %241 = "mhlo.broadcast_in_dim"(%139) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %242 = mhlo.add %238, %241 : tensor<2xui32>
    %243 = mhlo.add %242, %12 : tensor<2xui32>
    %244 = mhlo.add %240, %243 : tensor<2xui32>
    %245 = mhlo.shift_left %243, %11 : tensor<2xui32>
    %246 = mhlo.shift_right_logical %243, %10 : tensor<2xui32>
    %247 = mhlo.or %245, %246 : tensor<2xui32>
    %248 = mhlo.xor %244, %247 : tensor<2xui32>
    %249 = mhlo.add %244, %248 : tensor<2xui32>
    %250 = mhlo.shift_left %248, %9 : tensor<2xui32>
    %251 = mhlo.shift_right_logical %248, %8 : tensor<2xui32>
    %252 = mhlo.or %250, %251 : tensor<2xui32>
    %253 = mhlo.xor %249, %252 : tensor<2xui32>
    %254 = mhlo.add %249, %253 : tensor<2xui32>
    %255 = mhlo.shift_left %253, %6 : tensor<2xui32>
    %256 = mhlo.shift_right_logical %253, %7 : tensor<2xui32>
    %257 = mhlo.or %255, %256 : tensor<2xui32>
    %258 = mhlo.xor %254, %257 : tensor<2xui32>
    %259 = mhlo.add %254, %258 : tensor<2xui32>
    %260 = mhlo.shift_left %258, %7 : tensor<2xui32>
    %261 = mhlo.shift_right_logical %258, %6 : tensor<2xui32>
    %262 = mhlo.or %260, %261 : tensor<2xui32>
    %263 = mhlo.xor %259, %262 : tensor<2xui32>
    %264 = "mhlo.broadcast_in_dim"(%139) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %265 = mhlo.add %259, %264 : tensor<2xui32>
    %266 = "mhlo.broadcast_in_dim"(%133) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %267 = mhlo.add %263, %266 : tensor<2xui32>
    %268 = mhlo.add %267, %5 : tensor<2xui32>
    %269 = "mhlo.concatenate"(%265, %268) {dimension = 0 : i64} : (tensor<2xui32>, tensor<2xui32>) -> tensor<4xui32>
    %270 = "mhlo.reshape"(%269) : (tensor<4xui32>) -> tensor<2x2xui32>
    %271 = "mhlo.slice"(%270) {limit_indices = dense<[1, 2]> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<2x2xui32>) -> tensor<1x2xui32>
    %272 = "mhlo.reshape"(%271) : (tensor<1x2xui32>) -> tensor<2xui32>
    %273 = "mhlo.slice"(%270) {limit_indices = dense<2> : tensor<2xi64>, start_indices = dense<[1, 0]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<2x2xui32>) -> tensor<1x2xui32>
    %274 = "mhlo.reshape"(%273) : (tensor<1x2xui32>) -> tensor<2xui32>
    %275 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<4xui32>
    %276 = "mhlo.slice"(%274) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xui32>) -> tensor<1xui32>
    %277 = "mhlo.reshape"(%276) : (tensor<1xui32>) -> tensor<ui32>
    %278 = "mhlo.slice"(%274) {limit_indices = dense<2> : tensor<1xi64>, start_indices = dense<1> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xui32>) -> tensor<1xui32>
    %279 = "mhlo.reshape"(%278) : (tensor<1xui32>) -> tensor<ui32>
    %280 = "mhlo.slice"(%275) {limit_indices = dense<2> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<4xui32>) -> tensor<2xui32>
    %281 = "mhlo.slice"(%275) {limit_indices = dense<4> : tensor<1xi64>, start_indices = dense<2> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<4xui32>) -> tensor<2xui32>
    %282 = mhlo.xor %277, %279 : tensor<ui32>
    %283 = mhlo.xor %282, %20 : tensor<ui32>
    %284 = "mhlo.broadcast_in_dim"(%277) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %285 = mhlo.add %280, %284 : tensor<2xui32>
    %286 = "mhlo.broadcast_in_dim"(%279) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %287 = mhlo.add %281, %286 : tensor<2xui32>
    %288 = mhlo.add %285, %287 : tensor<2xui32>
    %289 = mhlo.shift_left %287, %11 : tensor<2xui32>
    %290 = mhlo.shift_right_logical %287, %10 : tensor<2xui32>
    %291 = mhlo.or %289, %290 : tensor<2xui32>
    %292 = mhlo.xor %288, %291 : tensor<2xui32>
    %293 = mhlo.add %288, %292 : tensor<2xui32>
    %294 = mhlo.shift_left %292, %9 : tensor<2xui32>
    %295 = mhlo.shift_right_logical %292, %8 : tensor<2xui32>
    %296 = mhlo.or %294, %295 : tensor<2xui32>
    %297 = mhlo.xor %293, %296 : tensor<2xui32>
    %298 = mhlo.add %293, %297 : tensor<2xui32>
    %299 = mhlo.shift_left %297, %6 : tensor<2xui32>
    %300 = mhlo.shift_right_logical %297, %7 : tensor<2xui32>
    %301 = mhlo.or %299, %300 : tensor<2xui32>
    %302 = mhlo.xor %298, %301 : tensor<2xui32>
    %303 = mhlo.add %298, %302 : tensor<2xui32>
    %304 = mhlo.shift_left %302, %7 : tensor<2xui32>
    %305 = mhlo.shift_right_logical %302, %6 : tensor<2xui32>
    %306 = mhlo.or %304, %305 : tensor<2xui32>
    %307 = mhlo.xor %303, %306 : tensor<2xui32>
    %308 = "mhlo.broadcast_in_dim"(%279) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %309 = mhlo.add %303, %308 : tensor<2xui32>
    %310 = "mhlo.broadcast_in_dim"(%283) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %311 = mhlo.add %307, %310 : tensor<2xui32>
    %312 = mhlo.add %311, %19 : tensor<2xui32>
    %313 = mhlo.add %309, %312 : tensor<2xui32>
    %314 = mhlo.shift_left %312, %8 : tensor<2xui32>
    %315 = mhlo.shift_right_logical %312, %9 : tensor<2xui32>
    %316 = mhlo.or %314, %315 : tensor<2xui32>
    %317 = mhlo.xor %313, %316 : tensor<2xui32>
    %318 = mhlo.add %313, %317 : tensor<2xui32>
    %319 = mhlo.shift_left %317, %17 : tensor<2xui32>
    %320 = mhlo.shift_right_logical %317, %16 : tensor<2xui32>
    %321 = mhlo.or %319, %320 : tensor<2xui32>
    %322 = mhlo.xor %318, %321 : tensor<2xui32>
    %323 = mhlo.add %318, %322 : tensor<2xui32>
    %324 = mhlo.shift_left %322, %15 : tensor<2xui32>
    %325 = mhlo.shift_right_logical %322, %15 : tensor<2xui32>
    %326 = mhlo.or %324, %325 : tensor<2xui32>
    %327 = mhlo.xor %323, %326 : tensor<2xui32>
    %328 = mhlo.add %323, %327 : tensor<2xui32>
    %329 = mhlo.shift_left %327, %14 : tensor<2xui32>
    %330 = mhlo.shift_right_logical %327, %13 : tensor<2xui32>
    %331 = mhlo.or %329, %330 : tensor<2xui32>
    %332 = mhlo.xor %328, %331 : tensor<2xui32>
    %333 = "mhlo.broadcast_in_dim"(%283) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %334 = mhlo.add %328, %333 : tensor<2xui32>
    %335 = "mhlo.broadcast_in_dim"(%277) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %336 = mhlo.add %332, %335 : tensor<2xui32>
    %337 = mhlo.add %336, %18 : tensor<2xui32>
    %338 = mhlo.add %334, %337 : tensor<2xui32>
    %339 = mhlo.shift_left %337, %11 : tensor<2xui32>
    %340 = mhlo.shift_right_logical %337, %10 : tensor<2xui32>
    %341 = mhlo.or %339, %340 : tensor<2xui32>
    %342 = mhlo.xor %338, %341 : tensor<2xui32>
    %343 = mhlo.add %338, %342 : tensor<2xui32>
    %344 = mhlo.shift_left %342, %9 : tensor<2xui32>
    %345 = mhlo.shift_right_logical %342, %8 : tensor<2xui32>
    %346 = mhlo.or %344, %345 : tensor<2xui32>
    %347 = mhlo.xor %343, %346 : tensor<2xui32>
    %348 = mhlo.add %343, %347 : tensor<2xui32>
    %349 = mhlo.shift_left %347, %6 : tensor<2xui32>
    %350 = mhlo.shift_right_logical %347, %7 : tensor<2xui32>
    %351 = mhlo.or %349, %350 : tensor<2xui32>
    %352 = mhlo.xor %348, %351 : tensor<2xui32>
    %353 = mhlo.add %348, %352 : tensor<2xui32>
    %354 = mhlo.shift_left %352, %7 : tensor<2xui32>
    %355 = mhlo.shift_right_logical %352, %6 : tensor<2xui32>
    %356 = mhlo.or %354, %355 : tensor<2xui32>
    %357 = mhlo.xor %353, %356 : tensor<2xui32>
    %358 = "mhlo.broadcast_in_dim"(%277) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %359 = mhlo.add %353, %358 : tensor<2xui32>
    %360 = "mhlo.broadcast_in_dim"(%279) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %361 = mhlo.add %357, %360 : tensor<2xui32>
    %362 = mhlo.add %361, %16 : tensor<2xui32>
    %363 = mhlo.add %359, %362 : tensor<2xui32>
    %364 = mhlo.shift_left %362, %8 : tensor<2xui32>
    %365 = mhlo.shift_right_logical %362, %9 : tensor<2xui32>
    %366 = mhlo.or %364, %365 : tensor<2xui32>
    %367 = mhlo.xor %363, %366 : tensor<2xui32>
    %368 = mhlo.add %363, %367 : tensor<2xui32>
    %369 = mhlo.shift_left %367, %17 : tensor<2xui32>
    %370 = mhlo.shift_right_logical %367, %16 : tensor<2xui32>
    %371 = mhlo.or %369, %370 : tensor<2xui32>
    %372 = mhlo.xor %368, %371 : tensor<2xui32>
    %373 = mhlo.add %368, %372 : tensor<2xui32>
    %374 = mhlo.shift_left %372, %15 : tensor<2xui32>
    %375 = mhlo.shift_right_logical %372, %15 : tensor<2xui32>
    %376 = mhlo.or %374, %375 : tensor<2xui32>
    %377 = mhlo.xor %373, %376 : tensor<2xui32>
    %378 = mhlo.add %373, %377 : tensor<2xui32>
    %379 = mhlo.shift_left %377, %14 : tensor<2xui32>
    %380 = mhlo.shift_right_logical %377, %13 : tensor<2xui32>
    %381 = mhlo.or %379, %380 : tensor<2xui32>
    %382 = mhlo.xor %378, %381 : tensor<2xui32>
    %383 = "mhlo.broadcast_in_dim"(%279) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %384 = mhlo.add %378, %383 : tensor<2xui32>
    %385 = "mhlo.broadcast_in_dim"(%283) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %386 = mhlo.add %382, %385 : tensor<2xui32>
    %387 = mhlo.add %386, %12 : tensor<2xui32>
    %388 = mhlo.add %384, %387 : tensor<2xui32>
    %389 = mhlo.shift_left %387, %11 : tensor<2xui32>
    %390 = mhlo.shift_right_logical %387, %10 : tensor<2xui32>
    %391 = mhlo.or %389, %390 : tensor<2xui32>
    %392 = mhlo.xor %388, %391 : tensor<2xui32>
    %393 = mhlo.add %388, %392 : tensor<2xui32>
    %394 = mhlo.shift_left %392, %9 : tensor<2xui32>
    %395 = mhlo.shift_right_logical %392, %8 : tensor<2xui32>
    %396 = mhlo.or %394, %395 : tensor<2xui32>
    %397 = mhlo.xor %393, %396 : tensor<2xui32>
    %398 = mhlo.add %393, %397 : tensor<2xui32>
    %399 = mhlo.shift_left %397, %6 : tensor<2xui32>
    %400 = mhlo.shift_right_logical %397, %7 : tensor<2xui32>
    %401 = mhlo.or %399, %400 : tensor<2xui32>
    %402 = mhlo.xor %398, %401 : tensor<2xui32>
    %403 = mhlo.add %398, %402 : tensor<2xui32>
    %404 = mhlo.shift_left %402, %7 : tensor<2xui32>
    %405 = mhlo.shift_right_logical %402, %6 : tensor<2xui32>
    %406 = mhlo.or %404, %405 : tensor<2xui32>
    %407 = mhlo.xor %403, %406 : tensor<2xui32>
    %408 = "mhlo.broadcast_in_dim"(%283) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %409 = mhlo.add %403, %408 : tensor<2xui32>
    %410 = "mhlo.broadcast_in_dim"(%277) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %411 = mhlo.add %407, %410 : tensor<2xui32>
    %412 = mhlo.add %411, %5 : tensor<2xui32>
    %413 = "mhlo.concatenate"(%409, %412) {dimension = 0 : i64} : (tensor<2xui32>, tensor<2xui32>) -> tensor<4xui32>
    %414 = "mhlo.reshape"(%413) : (tensor<4xui32>) -> tensor<2x2xui32>
    %415 = "mhlo.slice"(%414) {limit_indices = dense<[1, 2]> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<2x2xui32>) -> tensor<1x2xui32>
    %416 = "mhlo.reshape"(%415) : (tensor<1x2xui32>) -> tensor<2xui32>
    %417 = "mhlo.slice"(%414) {limit_indices = dense<2> : tensor<2xi64>, start_indices = dense<[1, 0]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<2x2xui32>) -> tensor<1x2xui32>
    %418 = "mhlo.reshape"(%417) : (tensor<1x2xui32>) -> tensor<2xui32>
    %419 = call @main_0(%68) : (tensor<f32>) -> tensor<f32>
    %420 = call @main_1(%23) : (tensor<f32>) -> tensor<f32>
    %421 = "mhlo.broadcast_in_dim"(%419) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1x1xf32>
    %422 = "mhlo.broadcast_in_dim"(%420) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1x1xf32>
    %423 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<802816xui32>
    %424 = "mhlo.slice"(%416) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xui32>) -> tensor<1xui32>
    %425 = "mhlo.reshape"(%424) : (tensor<1xui32>) -> tensor<ui32>
    %426 = "mhlo.slice"(%416) {limit_indices = dense<2> : tensor<1xi64>, start_indices = dense<1> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xui32>) -> tensor<1xui32>
    %427 = "mhlo.reshape"(%426) : (tensor<1xui32>) -> tensor<ui32>
    %428 = "mhlo.slice"(%423) {limit_indices = dense<401408> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<802816xui32>) -> tensor<401408xui32>
    %429 = "mhlo.slice"(%423) {limit_indices = dense<802816> : tensor<1xi64>, start_indices = dense<401408> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<802816xui32>) -> tensor<401408xui32>
    %430 = mhlo.xor %425, %427 : tensor<ui32>
    %431 = mhlo.xor %430, %20 : tensor<ui32>
    %432 = "mhlo.broadcast_in_dim"(%425) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<401408xui32>
    %433 = mhlo.add %428, %432 : tensor<401408xui32>
    %434 = "mhlo.broadcast_in_dim"(%427) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<401408xui32>
    %435 = mhlo.add %429, %434 : tensor<401408xui32>
    %436 = mhlo.add %433, %435 : tensor<401408xui32>
    %437 = mhlo.shift_left %435, %122 : tensor<401408xui32>
    %438 = mhlo.shift_right_logical %435, %121 : tensor<401408xui32>
    %439 = mhlo.or %437, %438 : tensor<401408xui32>
    %440 = mhlo.xor %436, %439 : tensor<401408xui32>
    %441 = mhlo.add %436, %440 : tensor<401408xui32>
    %442 = mhlo.shift_left %440, %120 : tensor<401408xui32>
    %443 = mhlo.shift_right_logical %440, %119 : tensor<401408xui32>
    %444 = mhlo.or %442, %443 : tensor<401408xui32>
    %445 = mhlo.xor %441, %444 : tensor<401408xui32>
    %446 = mhlo.add %441, %445 : tensor<401408xui32>
    %447 = mhlo.shift_left %445, %117 : tensor<401408xui32>
    %448 = mhlo.shift_right_logical %445, %118 : tensor<401408xui32>
    %449 = mhlo.or %447, %448 : tensor<401408xui32>
    %450 = mhlo.xor %446, %449 : tensor<401408xui32>
    %451 = mhlo.add %446, %450 : tensor<401408xui32>
    %452 = mhlo.shift_left %450, %118 : tensor<401408xui32>
    %453 = mhlo.shift_right_logical %450, %117 : tensor<401408xui32>
    %454 = mhlo.or %452, %453 : tensor<401408xui32>
    %455 = mhlo.xor %451, %454 : tensor<401408xui32>
    %456 = "mhlo.broadcast_in_dim"(%427) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<401408xui32>
    %457 = mhlo.add %451, %456 : tensor<401408xui32>
    %458 = "mhlo.broadcast_in_dim"(%431) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<401408xui32>
    %459 = mhlo.add %455, %458 : tensor<401408xui32>
    %460 = mhlo.add %459, %130 : tensor<401408xui32>
    %461 = mhlo.add %457, %460 : tensor<401408xui32>
    %462 = mhlo.shift_left %460, %119 : tensor<401408xui32>
    %463 = mhlo.shift_right_logical %460, %120 : tensor<401408xui32>
    %464 = mhlo.or %462, %463 : tensor<401408xui32>
    %465 = mhlo.xor %461, %464 : tensor<401408xui32>
    %466 = mhlo.add %461, %465 : tensor<401408xui32>
    %467 = mhlo.shift_left %465, %128 : tensor<401408xui32>
    %468 = mhlo.shift_right_logical %465, %127 : tensor<401408xui32>
    %469 = mhlo.or %467, %468 : tensor<401408xui32>
    %470 = mhlo.xor %466, %469 : tensor<401408xui32>
    %471 = mhlo.add %466, %470 : tensor<401408xui32>
    %472 = mhlo.shift_left %470, %126 : tensor<401408xui32>
    %473 = mhlo.shift_right_logical %470, %126 : tensor<401408xui32>
    %474 = mhlo.or %472, %473 : tensor<401408xui32>
    %475 = mhlo.xor %471, %474 : tensor<401408xui32>
    %476 = mhlo.add %471, %475 : tensor<401408xui32>
    %477 = mhlo.shift_left %475, %125 : tensor<401408xui32>
    %478 = mhlo.shift_right_logical %475, %124 : tensor<401408xui32>
    %479 = mhlo.or %477, %478 : tensor<401408xui32>
    %480 = mhlo.xor %476, %479 : tensor<401408xui32>
    %481 = "mhlo.broadcast_in_dim"(%431) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<401408xui32>
    %482 = mhlo.add %476, %481 : tensor<401408xui32>
    %483 = "mhlo.broadcast_in_dim"(%425) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<401408xui32>
    %484 = mhlo.add %480, %483 : tensor<401408xui32>
    %485 = mhlo.add %484, %129 : tensor<401408xui32>
    %486 = mhlo.add %482, %485 : tensor<401408xui32>
    %487 = mhlo.shift_left %485, %122 : tensor<401408xui32>
    %488 = mhlo.shift_right_logical %485, %121 : tensor<401408xui32>
    %489 = mhlo.or %487, %488 : tensor<401408xui32>
    %490 = mhlo.xor %486, %489 : tensor<401408xui32>
    %491 = mhlo.add %486, %490 : tensor<401408xui32>
    %492 = mhlo.shift_left %490, %120 : tensor<401408xui32>
    %493 = mhlo.shift_right_logical %490, %119 : tensor<401408xui32>
    %494 = mhlo.or %492, %493 : tensor<401408xui32>
    %495 = mhlo.xor %491, %494 : tensor<401408xui32>
    %496 = mhlo.add %491, %495 : tensor<401408xui32>
    %497 = mhlo.shift_left %495, %117 : tensor<401408xui32>
    %498 = mhlo.shift_right_logical %495, %118 : tensor<401408xui32>
    %499 = mhlo.or %497, %498 : tensor<401408xui32>
    %500 = mhlo.xor %496, %499 : tensor<401408xui32>
    %501 = mhlo.add %496, %500 : tensor<401408xui32>
    %502 = mhlo.shift_left %500, %118 : tensor<401408xui32>
    %503 = mhlo.shift_right_logical %500, %117 : tensor<401408xui32>
    %504 = mhlo.or %502, %503 : tensor<401408xui32>
    %505 = mhlo.xor %501, %504 : tensor<401408xui32>
    %506 = "mhlo.broadcast_in_dim"(%425) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<401408xui32>
    %507 = mhlo.add %501, %506 : tensor<401408xui32>
    %508 = "mhlo.broadcast_in_dim"(%427) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<401408xui32>
    %509 = mhlo.add %505, %508 : tensor<401408xui32>
    %510 = mhlo.add %509, %127 : tensor<401408xui32>
    %511 = mhlo.add %507, %510 : tensor<401408xui32>
    %512 = mhlo.shift_left %510, %119 : tensor<401408xui32>
    %513 = mhlo.shift_right_logical %510, %120 : tensor<401408xui32>
    %514 = mhlo.or %512, %513 : tensor<401408xui32>
    %515 = mhlo.xor %511, %514 : tensor<401408xui32>
    %516 = mhlo.add %511, %515 : tensor<401408xui32>
    %517 = mhlo.shift_left %515, %128 : tensor<401408xui32>
    %518 = mhlo.shift_right_logical %515, %127 : tensor<401408xui32>
    %519 = mhlo.or %517, %518 : tensor<401408xui32>
    %520 = mhlo.xor %516, %519 : tensor<401408xui32>
    %521 = mhlo.add %516, %520 : tensor<401408xui32>
    %522 = mhlo.shift_left %520, %126 : tensor<401408xui32>
    %523 = mhlo.shift_right_logical %520, %126 : tensor<401408xui32>
    %524 = mhlo.or %522, %523 : tensor<401408xui32>
    %525 = mhlo.xor %521, %524 : tensor<401408xui32>
    %526 = mhlo.add %521, %525 : tensor<401408xui32>
    %527 = mhlo.shift_left %525, %125 : tensor<401408xui32>
    %528 = mhlo.shift_right_logical %525, %124 : tensor<401408xui32>
    %529 = mhlo.or %527, %528 : tensor<401408xui32>
    %530 = mhlo.xor %526, %529 : tensor<401408xui32>
    %531 = "mhlo.broadcast_in_dim"(%427) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<401408xui32>
    %532 = mhlo.add %526, %531 : tensor<401408xui32>
    %533 = "mhlo.broadcast_in_dim"(%431) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<401408xui32>
    %534 = mhlo.add %530, %533 : tensor<401408xui32>
    %535 = mhlo.add %534, %123 : tensor<401408xui32>
    %536 = mhlo.add %532, %535 : tensor<401408xui32>
    %537 = mhlo.shift_left %535, %122 : tensor<401408xui32>
    %538 = mhlo.shift_right_logical %535, %121 : tensor<401408xui32>
    %539 = mhlo.or %537, %538 : tensor<401408xui32>
    %540 = mhlo.xor %536, %539 : tensor<401408xui32>
    %541 = mhlo.add %536, %540 : tensor<401408xui32>
    %542 = mhlo.shift_left %540, %120 : tensor<401408xui32>
    %543 = mhlo.shift_right_logical %540, %119 : tensor<401408xui32>
    %544 = mhlo.or %542, %543 : tensor<401408xui32>
    %545 = mhlo.xor %541, %544 : tensor<401408xui32>
    %546 = mhlo.add %541, %545 : tensor<401408xui32>
    %547 = mhlo.shift_left %545, %117 : tensor<401408xui32>
    %548 = mhlo.shift_right_logical %545, %118 : tensor<401408xui32>
    %549 = mhlo.or %547, %548 : tensor<401408xui32>
    %550 = mhlo.xor %546, %549 : tensor<401408xui32>
    %551 = mhlo.add %546, %550 : tensor<401408xui32>
    %552 = mhlo.shift_left %550, %118 : tensor<401408xui32>
    %553 = mhlo.shift_right_logical %550, %117 : tensor<401408xui32>
    %554 = mhlo.or %552, %553 : tensor<401408xui32>
    %555 = mhlo.xor %551, %554 : tensor<401408xui32>
    %556 = "mhlo.broadcast_in_dim"(%431) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<401408xui32>
    %557 = mhlo.add %551, %556 : tensor<401408xui32>
    %558 = "mhlo.broadcast_in_dim"(%425) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<401408xui32>
    %559 = mhlo.add %555, %558 : tensor<401408xui32>
    %560 = mhlo.add %559, %116 : tensor<401408xui32>
    %561 = "mhlo.concatenate"(%557, %560) {dimension = 0 : i64} : (tensor<401408xui32>, tensor<401408xui32>) -> tensor<802816xui32>
    %562 = "mhlo.reshape"(%561) : (tensor<802816xui32>) -> tensor<784x1024xui32>
    %563 = mhlo.shift_right_logical %562, %115 : tensor<784x1024xui32>
    %564 = mhlo.or %563, %114 : tensor<784x1024xui32>
    %565 = "mhlo.bitcast_convert"(%564) : (tensor<784x1024xui32>) -> tensor<784x1024xf32>
    %566 = mhlo.subtract %565, %113 : tensor<784x1024xf32>
    %567 = mhlo.subtract %422, %421 : tensor<1x1xf32>
    %568 = "mhlo.broadcast_in_dim"(%567) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor<1x1xf32>) -> tensor<784x1024xf32>
    %569 = mhlo.multiply %566, %568 : tensor<784x1024xf32>
    %570 = "mhlo.broadcast_in_dim"(%421) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor<1x1xf32>) -> tensor<784x1024xf32>
    %571 = mhlo.add %569, %570 : tensor<784x1024xf32>
    %572 = "mhlo.broadcast_in_dim"(%421) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor<1x1xf32>) -> tensor<784x1024xf32>
    %573 = mhlo.maximum %572, %571 : tensor<784x1024xf32>
    %574 = call @main_2(%573) : (tensor<784x1024xf32>) -> tensor<784x1024xf32>
    %575 = mhlo.multiply %574, %112 : tensor<784x1024xf32>
    %576 = chlo.next_after %48, %47 : tensor<f32>, tensor<f32> -> tensor<f32>
    %577 = chlo.next_after %46, %45 : tensor<f32>, tensor<f32> -> tensor<f32>
    %578 = call @jit_clip(%575, %576, %577) : (tensor<784x1024xf32>, tensor<f32>, tensor<f32>) -> tensor<784x1024xf32>
    %579 = mhlo.multiply %578, %111 : tensor<784x1024xf32>
    %580 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<1024xui32>
    %581 = "mhlo.slice"(%418) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xui32>) -> tensor<1xui32>
    %582 = "mhlo.reshape"(%581) : (tensor<1xui32>) -> tensor<ui32>
    %583 = "mhlo.slice"(%418) {limit_indices = dense<2> : tensor<1xi64>, start_indices = dense<1> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xui32>) -> tensor<1xui32>
    %584 = "mhlo.reshape"(%583) : (tensor<1xui32>) -> tensor<ui32>
    %585 = "mhlo.slice"(%580) {limit_indices = dense<512> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1024xui32>) -> tensor<512xui32>
    %586 = "mhlo.slice"(%580) {limit_indices = dense<1024> : tensor<1xi64>, start_indices = dense<512> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1024xui32>) -> tensor<512xui32>
    %587 = mhlo.xor %582, %584 : tensor<ui32>
    %588 = mhlo.xor %587, %20 : tensor<ui32>
    %589 = "mhlo.broadcast_in_dim"(%582) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<512xui32>
    %590 = mhlo.add %585, %589 : tensor<512xui32>
    %591 = "mhlo.broadcast_in_dim"(%584) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<512xui32>
    %592 = mhlo.add %586, %591 : tensor<512xui32>
    %593 = mhlo.add %590, %592 : tensor<512xui32>
    %594 = mhlo.shift_left %592, %82 : tensor<512xui32>
    %595 = mhlo.shift_right_logical %592, %81 : tensor<512xui32>
    %596 = mhlo.or %594, %595 : tensor<512xui32>
    %597 = mhlo.xor %593, %596 : tensor<512xui32>
    %598 = mhlo.add %593, %597 : tensor<512xui32>
    %599 = mhlo.shift_left %597, %80 : tensor<512xui32>
    %600 = mhlo.shift_right_logical %597, %79 : tensor<512xui32>
    %601 = mhlo.or %599, %600 : tensor<512xui32>
    %602 = mhlo.xor %598, %601 : tensor<512xui32>
    %603 = mhlo.add %598, %602 : tensor<512xui32>
    %604 = mhlo.shift_left %602, %77 : tensor<512xui32>
    %605 = mhlo.shift_right_logical %602, %78 : tensor<512xui32>
    %606 = mhlo.or %604, %605 : tensor<512xui32>
    %607 = mhlo.xor %603, %606 : tensor<512xui32>
    %608 = mhlo.add %603, %607 : tensor<512xui32>
    %609 = mhlo.shift_left %607, %78 : tensor<512xui32>
    %610 = mhlo.shift_right_logical %607, %77 : tensor<512xui32>
    %611 = mhlo.or %609, %610 : tensor<512xui32>
    %612 = mhlo.xor %608, %611 : tensor<512xui32>
    %613 = "mhlo.broadcast_in_dim"(%584) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<512xui32>
    %614 = mhlo.add %608, %613 : tensor<512xui32>
    %615 = "mhlo.broadcast_in_dim"(%588) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<512xui32>
    %616 = mhlo.add %612, %615 : tensor<512xui32>
    %617 = mhlo.add %616, %90 : tensor<512xui32>
    %618 = mhlo.add %614, %617 : tensor<512xui32>
    %619 = mhlo.shift_left %617, %79 : tensor<512xui32>
    %620 = mhlo.shift_right_logical %617, %80 : tensor<512xui32>
    %621 = mhlo.or %619, %620 : tensor<512xui32>
    %622 = mhlo.xor %618, %621 : tensor<512xui32>
    %623 = mhlo.add %618, %622 : tensor<512xui32>
    %624 = mhlo.shift_left %622, %88 : tensor<512xui32>
    %625 = mhlo.shift_right_logical %622, %87 : tensor<512xui32>
    %626 = mhlo.or %624, %625 : tensor<512xui32>
    %627 = mhlo.xor %623, %626 : tensor<512xui32>
    %628 = mhlo.add %623, %627 : tensor<512xui32>
    %629 = mhlo.shift_left %627, %86 : tensor<512xui32>
    %630 = mhlo.shift_right_logical %627, %86 : tensor<512xui32>
    %631 = mhlo.or %629, %630 : tensor<512xui32>
    %632 = mhlo.xor %628, %631 : tensor<512xui32>
    %633 = mhlo.add %628, %632 : tensor<512xui32>
    %634 = mhlo.shift_left %632, %85 : tensor<512xui32>
    %635 = mhlo.shift_right_logical %632, %84 : tensor<512xui32>
    %636 = mhlo.or %634, %635 : tensor<512xui32>
    %637 = mhlo.xor %633, %636 : tensor<512xui32>
    %638 = "mhlo.broadcast_in_dim"(%588) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<512xui32>
    %639 = mhlo.add %633, %638 : tensor<512xui32>
    %640 = "mhlo.broadcast_in_dim"(%582) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<512xui32>
    %641 = mhlo.add %637, %640 : tensor<512xui32>
    %642 = mhlo.add %641, %89 : tensor<512xui32>
    %643 = mhlo.add %639, %642 : tensor<512xui32>
    %644 = mhlo.shift_left %642, %82 : tensor<512xui32>
    %645 = mhlo.shift_right_logical %642, %81 : tensor<512xui32>
    %646 = mhlo.or %644, %645 : tensor<512xui32>
    %647 = mhlo.xor %643, %646 : tensor<512xui32>
    %648 = mhlo.add %643, %647 : tensor<512xui32>
    %649 = mhlo.shift_left %647, %80 : tensor<512xui32>
    %650 = mhlo.shift_right_logical %647, %79 : tensor<512xui32>
    %651 = mhlo.or %649, %650 : tensor<512xui32>
    %652 = mhlo.xor %648, %651 : tensor<512xui32>
    %653 = mhlo.add %648, %652 : tensor<512xui32>
    %654 = mhlo.shift_left %652, %77 : tensor<512xui32>
    %655 = mhlo.shift_right_logical %652, %78 : tensor<512xui32>
    %656 = mhlo.or %654, %655 : tensor<512xui32>
    %657 = mhlo.xor %653, %656 : tensor<512xui32>
    %658 = mhlo.add %653, %657 : tensor<512xui32>
    %659 = mhlo.shift_left %657, %78 : tensor<512xui32>
    %660 = mhlo.shift_right_logical %657, %77 : tensor<512xui32>
    %661 = mhlo.or %659, %660 : tensor<512xui32>
    %662 = mhlo.xor %658, %661 : tensor<512xui32>
    %663 = "mhlo.broadcast_in_dim"(%582) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<512xui32>
    %664 = mhlo.add %658, %663 : tensor<512xui32>
    %665 = "mhlo.broadcast_in_dim"(%584) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<512xui32>
    %666 = mhlo.add %662, %665 : tensor<512xui32>
    %667 = mhlo.add %666, %87 : tensor<512xui32>
    %668 = mhlo.add %664, %667 : tensor<512xui32>
    %669 = mhlo.shift_left %667, %79 : tensor<512xui32>
    %670 = mhlo.shift_right_logical %667, %80 : tensor<512xui32>
    %671 = mhlo.or %669, %670 : tensor<512xui32>
    %672 = mhlo.xor %668, %671 : tensor<512xui32>
    %673 = mhlo.add %668, %672 : tensor<512xui32>
    %674 = mhlo.shift_left %672, %88 : tensor<512xui32>
    %675 = mhlo.shift_right_logical %672, %87 : tensor<512xui32>
    %676 = mhlo.or %674, %675 : tensor<512xui32>
    %677 = mhlo.xor %673, %676 : tensor<512xui32>
    %678 = mhlo.add %673, %677 : tensor<512xui32>
    %679 = mhlo.shift_left %677, %86 : tensor<512xui32>
    %680 = mhlo.shift_right_logical %677, %86 : tensor<512xui32>
    %681 = mhlo.or %679, %680 : tensor<512xui32>
    %682 = mhlo.xor %678, %681 : tensor<512xui32>
    %683 = mhlo.add %678, %682 : tensor<512xui32>
    %684 = mhlo.shift_left %682, %85 : tensor<512xui32>
    %685 = mhlo.shift_right_logical %682, %84 : tensor<512xui32>
    %686 = mhlo.or %684, %685 : tensor<512xui32>
    %687 = mhlo.xor %683, %686 : tensor<512xui32>
    %688 = "mhlo.broadcast_in_dim"(%584) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<512xui32>
    %689 = mhlo.add %683, %688 : tensor<512xui32>
    %690 = "mhlo.broadcast_in_dim"(%588) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<512xui32>
    %691 = mhlo.add %687, %690 : tensor<512xui32>
    %692 = mhlo.add %691, %83 : tensor<512xui32>
    %693 = mhlo.add %689, %692 : tensor<512xui32>
    %694 = mhlo.shift_left %692, %82 : tensor<512xui32>
    %695 = mhlo.shift_right_logical %692, %81 : tensor<512xui32>
    %696 = mhlo.or %694, %695 : tensor<512xui32>
    %697 = mhlo.xor %693, %696 : tensor<512xui32>
    %698 = mhlo.add %693, %697 : tensor<512xui32>
    %699 = mhlo.shift_left %697, %80 : tensor<512xui32>
    %700 = mhlo.shift_right_logical %697, %79 : tensor<512xui32>
    %701 = mhlo.or %699, %700 : tensor<512xui32>
    %702 = mhlo.xor %698, %701 : tensor<512xui32>
    %703 = mhlo.add %698, %702 : tensor<512xui32>
    %704 = mhlo.shift_left %702, %77 : tensor<512xui32>
    %705 = mhlo.shift_right_logical %702, %78 : tensor<512xui32>
    %706 = mhlo.or %704, %705 : tensor<512xui32>
    %707 = mhlo.xor %703, %706 : tensor<512xui32>
    %708 = mhlo.add %703, %707 : tensor<512xui32>
    %709 = mhlo.shift_left %707, %78 : tensor<512xui32>
    %710 = mhlo.shift_right_logical %707, %77 : tensor<512xui32>
    %711 = mhlo.or %709, %710 : tensor<512xui32>
    %712 = mhlo.xor %708, %711 : tensor<512xui32>
    %713 = "mhlo.broadcast_in_dim"(%588) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<512xui32>
    %714 = mhlo.add %708, %713 : tensor<512xui32>
    %715 = "mhlo.broadcast_in_dim"(%582) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<512xui32>
    %716 = mhlo.add %712, %715 : tensor<512xui32>
    %717 = mhlo.add %716, %76 : tensor<512xui32>
    %718 = "mhlo.concatenate"(%714, %717) {dimension = 0 : i64} : (tensor<512xui32>, tensor<512xui32>) -> tensor<1024xui32>
    %719 = mhlo.shift_right_logical %718, %75 : tensor<1024xui32>
    %720 = mhlo.or %719, %74 : tensor<1024xui32>
    %721 = "mhlo.bitcast_convert"(%720) : (tensor<1024xui32>) -> tensor<1024xf32>
    %722 = mhlo.subtract %721, %73 : tensor<1024xf32>
    %723 = mhlo.multiply %722, %72 : tensor<1024xf32>
    %724 = mhlo.add %723, %71 : tensor<1024xf32>
    %725 = mhlo.maximum %724, %71 : tensor<1024xf32>
    %726 = call @main_3(%725) : (tensor<1024xf32>) -> tensor<1024xf32>
    %727 = mhlo.multiply %726, %70 : tensor<1024xf32>
    %728 = mhlo.multiply %727, %69 : tensor<1024xf32>
    %729 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<4xui32>
    %730 = "mhlo.slice"(%272) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xui32>) -> tensor<1xui32>
    %731 = "mhlo.reshape"(%730) : (tensor<1xui32>) -> tensor<ui32>
    %732 = "mhlo.slice"(%272) {limit_indices = dense<2> : tensor<1xi64>, start_indices = dense<1> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xui32>) -> tensor<1xui32>
    %733 = "mhlo.reshape"(%732) : (tensor<1xui32>) -> tensor<ui32>
    %734 = "mhlo.slice"(%729) {limit_indices = dense<2> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<4xui32>) -> tensor<2xui32>
    %735 = "mhlo.slice"(%729) {limit_indices = dense<4> : tensor<1xi64>, start_indices = dense<2> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<4xui32>) -> tensor<2xui32>
    %736 = mhlo.xor %731, %733 : tensor<ui32>
    %737 = mhlo.xor %736, %20 : tensor<ui32>
    %738 = "mhlo.broadcast_in_dim"(%731) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %739 = mhlo.add %734, %738 : tensor<2xui32>
    %740 = "mhlo.broadcast_in_dim"(%733) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %741 = mhlo.add %735, %740 : tensor<2xui32>
    %742 = mhlo.add %739, %741 : tensor<2xui32>
    %743 = mhlo.shift_left %741, %11 : tensor<2xui32>
    %744 = mhlo.shift_right_logical %741, %10 : tensor<2xui32>
    %745 = mhlo.or %743, %744 : tensor<2xui32>
    %746 = mhlo.xor %742, %745 : tensor<2xui32>
    %747 = mhlo.add %742, %746 : tensor<2xui32>
    %748 = mhlo.shift_left %746, %9 : tensor<2xui32>
    %749 = mhlo.shift_right_logical %746, %8 : tensor<2xui32>
    %750 = mhlo.or %748, %749 : tensor<2xui32>
    %751 = mhlo.xor %747, %750 : tensor<2xui32>
    %752 = mhlo.add %747, %751 : tensor<2xui32>
    %753 = mhlo.shift_left %751, %6 : tensor<2xui32>
    %754 = mhlo.shift_right_logical %751, %7 : tensor<2xui32>
    %755 = mhlo.or %753, %754 : tensor<2xui32>
    %756 = mhlo.xor %752, %755 : tensor<2xui32>
    %757 = mhlo.add %752, %756 : tensor<2xui32>
    %758 = mhlo.shift_left %756, %7 : tensor<2xui32>
    %759 = mhlo.shift_right_logical %756, %6 : tensor<2xui32>
    %760 = mhlo.or %758, %759 : tensor<2xui32>
    %761 = mhlo.xor %757, %760 : tensor<2xui32>
    %762 = "mhlo.broadcast_in_dim"(%733) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %763 = mhlo.add %757, %762 : tensor<2xui32>
    %764 = "mhlo.broadcast_in_dim"(%737) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %765 = mhlo.add %761, %764 : tensor<2xui32>
    %766 = mhlo.add %765, %19 : tensor<2xui32>
    %767 = mhlo.add %763, %766 : tensor<2xui32>
    %768 = mhlo.shift_left %766, %8 : tensor<2xui32>
    %769 = mhlo.shift_right_logical %766, %9 : tensor<2xui32>
    %770 = mhlo.or %768, %769 : tensor<2xui32>
    %771 = mhlo.xor %767, %770 : tensor<2xui32>
    %772 = mhlo.add %767, %771 : tensor<2xui32>
    %773 = mhlo.shift_left %771, %17 : tensor<2xui32>
    %774 = mhlo.shift_right_logical %771, %16 : tensor<2xui32>
    %775 = mhlo.or %773, %774 : tensor<2xui32>
    %776 = mhlo.xor %772, %775 : tensor<2xui32>
    %777 = mhlo.add %772, %776 : tensor<2xui32>
    %778 = mhlo.shift_left %776, %15 : tensor<2xui32>
    %779 = mhlo.shift_right_logical %776, %15 : tensor<2xui32>
    %780 = mhlo.or %778, %779 : tensor<2xui32>
    %781 = mhlo.xor %777, %780 : tensor<2xui32>
    %782 = mhlo.add %777, %781 : tensor<2xui32>
    %783 = mhlo.shift_left %781, %14 : tensor<2xui32>
    %784 = mhlo.shift_right_logical %781, %13 : tensor<2xui32>
    %785 = mhlo.or %783, %784 : tensor<2xui32>
    %786 = mhlo.xor %782, %785 : tensor<2xui32>
    %787 = "mhlo.broadcast_in_dim"(%737) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %788 = mhlo.add %782, %787 : tensor<2xui32>
    %789 = "mhlo.broadcast_in_dim"(%731) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %790 = mhlo.add %786, %789 : tensor<2xui32>
    %791 = mhlo.add %790, %18 : tensor<2xui32>
    %792 = mhlo.add %788, %791 : tensor<2xui32>
    %793 = mhlo.shift_left %791, %11 : tensor<2xui32>
    %794 = mhlo.shift_right_logical %791, %10 : tensor<2xui32>
    %795 = mhlo.or %793, %794 : tensor<2xui32>
    %796 = mhlo.xor %792, %795 : tensor<2xui32>
    %797 = mhlo.add %792, %796 : tensor<2xui32>
    %798 = mhlo.shift_left %796, %9 : tensor<2xui32>
    %799 = mhlo.shift_right_logical %796, %8 : tensor<2xui32>
    %800 = mhlo.or %798, %799 : tensor<2xui32>
    %801 = mhlo.xor %797, %800 : tensor<2xui32>
    %802 = mhlo.add %797, %801 : tensor<2xui32>
    %803 = mhlo.shift_left %801, %6 : tensor<2xui32>
    %804 = mhlo.shift_right_logical %801, %7 : tensor<2xui32>
    %805 = mhlo.or %803, %804 : tensor<2xui32>
    %806 = mhlo.xor %802, %805 : tensor<2xui32>
    %807 = mhlo.add %802, %806 : tensor<2xui32>
    %808 = mhlo.shift_left %806, %7 : tensor<2xui32>
    %809 = mhlo.shift_right_logical %806, %6 : tensor<2xui32>
    %810 = mhlo.or %808, %809 : tensor<2xui32>
    %811 = mhlo.xor %807, %810 : tensor<2xui32>
    %812 = "mhlo.broadcast_in_dim"(%731) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %813 = mhlo.add %807, %812 : tensor<2xui32>
    %814 = "mhlo.broadcast_in_dim"(%733) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %815 = mhlo.add %811, %814 : tensor<2xui32>
    %816 = mhlo.add %815, %16 : tensor<2xui32>
    %817 = mhlo.add %813, %816 : tensor<2xui32>
    %818 = mhlo.shift_left %816, %8 : tensor<2xui32>
    %819 = mhlo.shift_right_logical %816, %9 : tensor<2xui32>
    %820 = mhlo.or %818, %819 : tensor<2xui32>
    %821 = mhlo.xor %817, %820 : tensor<2xui32>
    %822 = mhlo.add %817, %821 : tensor<2xui32>
    %823 = mhlo.shift_left %821, %17 : tensor<2xui32>
    %824 = mhlo.shift_right_logical %821, %16 : tensor<2xui32>
    %825 = mhlo.or %823, %824 : tensor<2xui32>
    %826 = mhlo.xor %822, %825 : tensor<2xui32>
    %827 = mhlo.add %822, %826 : tensor<2xui32>
    %828 = mhlo.shift_left %826, %15 : tensor<2xui32>
    %829 = mhlo.shift_right_logical %826, %15 : tensor<2xui32>
    %830 = mhlo.or %828, %829 : tensor<2xui32>
    %831 = mhlo.xor %827, %830 : tensor<2xui32>
    %832 = mhlo.add %827, %831 : tensor<2xui32>
    %833 = mhlo.shift_left %831, %14 : tensor<2xui32>
    %834 = mhlo.shift_right_logical %831, %13 : tensor<2xui32>
    %835 = mhlo.or %833, %834 : tensor<2xui32>
    %836 = mhlo.xor %832, %835 : tensor<2xui32>
    %837 = "mhlo.broadcast_in_dim"(%733) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %838 = mhlo.add %832, %837 : tensor<2xui32>
    %839 = "mhlo.broadcast_in_dim"(%737) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %840 = mhlo.add %836, %839 : tensor<2xui32>
    %841 = mhlo.add %840, %12 : tensor<2xui32>
    %842 = mhlo.add %838, %841 : tensor<2xui32>
    %843 = mhlo.shift_left %841, %11 : tensor<2xui32>
    %844 = mhlo.shift_right_logical %841, %10 : tensor<2xui32>
    %845 = mhlo.or %843, %844 : tensor<2xui32>
    %846 = mhlo.xor %842, %845 : tensor<2xui32>
    %847 = mhlo.add %842, %846 : tensor<2xui32>
    %848 = mhlo.shift_left %846, %9 : tensor<2xui32>
    %849 = mhlo.shift_right_logical %846, %8 : tensor<2xui32>
    %850 = mhlo.or %848, %849 : tensor<2xui32>
    %851 = mhlo.xor %847, %850 : tensor<2xui32>
    %852 = mhlo.add %847, %851 : tensor<2xui32>
    %853 = mhlo.shift_left %851, %6 : tensor<2xui32>
    %854 = mhlo.shift_right_logical %851, %7 : tensor<2xui32>
    %855 = mhlo.or %853, %854 : tensor<2xui32>
    %856 = mhlo.xor %852, %855 : tensor<2xui32>
    %857 = mhlo.add %852, %856 : tensor<2xui32>
    %858 = mhlo.shift_left %856, %7 : tensor<2xui32>
    %859 = mhlo.shift_right_logical %856, %6 : tensor<2xui32>
    %860 = mhlo.or %858, %859 : tensor<2xui32>
    %861 = mhlo.xor %857, %860 : tensor<2xui32>
    %862 = "mhlo.broadcast_in_dim"(%737) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %863 = mhlo.add %857, %862 : tensor<2xui32>
    %864 = "mhlo.broadcast_in_dim"(%731) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %865 = mhlo.add %861, %864 : tensor<2xui32>
    %866 = mhlo.add %865, %5 : tensor<2xui32>
    %867 = "mhlo.concatenate"(%863, %866) {dimension = 0 : i64} : (tensor<2xui32>, tensor<2xui32>) -> tensor<4xui32>
    %868 = "mhlo.reshape"(%867) : (tensor<4xui32>) -> tensor<2x2xui32>
    %869 = "mhlo.slice"(%868) {limit_indices = dense<[1, 2]> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<2x2xui32>) -> tensor<1x2xui32>
    %870 = "mhlo.reshape"(%869) : (tensor<1x2xui32>) -> tensor<2xui32>
    %871 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<4xui32>
    %872 = "mhlo.slice"(%870) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xui32>) -> tensor<1xui32>
    %873 = "mhlo.reshape"(%872) : (tensor<1xui32>) -> tensor<ui32>
    %874 = "mhlo.slice"(%870) {limit_indices = dense<2> : tensor<1xi64>, start_indices = dense<1> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xui32>) -> tensor<1xui32>
    %875 = "mhlo.reshape"(%874) : (tensor<1xui32>) -> tensor<ui32>
    %876 = "mhlo.slice"(%871) {limit_indices = dense<2> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<4xui32>) -> tensor<2xui32>
    %877 = "mhlo.slice"(%871) {limit_indices = dense<4> : tensor<1xi64>, start_indices = dense<2> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<4xui32>) -> tensor<2xui32>
    %878 = mhlo.xor %873, %875 : tensor<ui32>
    %879 = mhlo.xor %878, %20 : tensor<ui32>
    %880 = "mhlo.broadcast_in_dim"(%873) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %881 = mhlo.add %876, %880 : tensor<2xui32>
    %882 = "mhlo.broadcast_in_dim"(%875) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %883 = mhlo.add %877, %882 : tensor<2xui32>
    %884 = mhlo.add %881, %883 : tensor<2xui32>
    %885 = mhlo.shift_left %883, %11 : tensor<2xui32>
    %886 = mhlo.shift_right_logical %883, %10 : tensor<2xui32>
    %887 = mhlo.or %885, %886 : tensor<2xui32>
    %888 = mhlo.xor %884, %887 : tensor<2xui32>
    %889 = mhlo.add %884, %888 : tensor<2xui32>
    %890 = mhlo.shift_left %888, %9 : tensor<2xui32>
    %891 = mhlo.shift_right_logical %888, %8 : tensor<2xui32>
    %892 = mhlo.or %890, %891 : tensor<2xui32>
    %893 = mhlo.xor %889, %892 : tensor<2xui32>
    %894 = mhlo.add %889, %893 : tensor<2xui32>
    %895 = mhlo.shift_left %893, %6 : tensor<2xui32>
    %896 = mhlo.shift_right_logical %893, %7 : tensor<2xui32>
    %897 = mhlo.or %895, %896 : tensor<2xui32>
    %898 = mhlo.xor %894, %897 : tensor<2xui32>
    %899 = mhlo.add %894, %898 : tensor<2xui32>
    %900 = mhlo.shift_left %898, %7 : tensor<2xui32>
    %901 = mhlo.shift_right_logical %898, %6 : tensor<2xui32>
    %902 = mhlo.or %900, %901 : tensor<2xui32>
    %903 = mhlo.xor %899, %902 : tensor<2xui32>
    %904 = "mhlo.broadcast_in_dim"(%875) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %905 = mhlo.add %899, %904 : tensor<2xui32>
    %906 = "mhlo.broadcast_in_dim"(%879) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %907 = mhlo.add %903, %906 : tensor<2xui32>
    %908 = mhlo.add %907, %19 : tensor<2xui32>
    %909 = mhlo.add %905, %908 : tensor<2xui32>
    %910 = mhlo.shift_left %908, %8 : tensor<2xui32>
    %911 = mhlo.shift_right_logical %908, %9 : tensor<2xui32>
    %912 = mhlo.or %910, %911 : tensor<2xui32>
    %913 = mhlo.xor %909, %912 : tensor<2xui32>
    %914 = mhlo.add %909, %913 : tensor<2xui32>
    %915 = mhlo.shift_left %913, %17 : tensor<2xui32>
    %916 = mhlo.shift_right_logical %913, %16 : tensor<2xui32>
    %917 = mhlo.or %915, %916 : tensor<2xui32>
    %918 = mhlo.xor %914, %917 : tensor<2xui32>
    %919 = mhlo.add %914, %918 : tensor<2xui32>
    %920 = mhlo.shift_left %918, %15 : tensor<2xui32>
    %921 = mhlo.shift_right_logical %918, %15 : tensor<2xui32>
    %922 = mhlo.or %920, %921 : tensor<2xui32>
    %923 = mhlo.xor %919, %922 : tensor<2xui32>
    %924 = mhlo.add %919, %923 : tensor<2xui32>
    %925 = mhlo.shift_left %923, %14 : tensor<2xui32>
    %926 = mhlo.shift_right_logical %923, %13 : tensor<2xui32>
    %927 = mhlo.or %925, %926 : tensor<2xui32>
    %928 = mhlo.xor %924, %927 : tensor<2xui32>
    %929 = "mhlo.broadcast_in_dim"(%879) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %930 = mhlo.add %924, %929 : tensor<2xui32>
    %931 = "mhlo.broadcast_in_dim"(%873) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %932 = mhlo.add %928, %931 : tensor<2xui32>
    %933 = mhlo.add %932, %18 : tensor<2xui32>
    %934 = mhlo.add %930, %933 : tensor<2xui32>
    %935 = mhlo.shift_left %933, %11 : tensor<2xui32>
    %936 = mhlo.shift_right_logical %933, %10 : tensor<2xui32>
    %937 = mhlo.or %935, %936 : tensor<2xui32>
    %938 = mhlo.xor %934, %937 : tensor<2xui32>
    %939 = mhlo.add %934, %938 : tensor<2xui32>
    %940 = mhlo.shift_left %938, %9 : tensor<2xui32>
    %941 = mhlo.shift_right_logical %938, %8 : tensor<2xui32>
    %942 = mhlo.or %940, %941 : tensor<2xui32>
    %943 = mhlo.xor %939, %942 : tensor<2xui32>
    %944 = mhlo.add %939, %943 : tensor<2xui32>
    %945 = mhlo.shift_left %943, %6 : tensor<2xui32>
    %946 = mhlo.shift_right_logical %943, %7 : tensor<2xui32>
    %947 = mhlo.or %945, %946 : tensor<2xui32>
    %948 = mhlo.xor %944, %947 : tensor<2xui32>
    %949 = mhlo.add %944, %948 : tensor<2xui32>
    %950 = mhlo.shift_left %948, %7 : tensor<2xui32>
    %951 = mhlo.shift_right_logical %948, %6 : tensor<2xui32>
    %952 = mhlo.or %950, %951 : tensor<2xui32>
    %953 = mhlo.xor %949, %952 : tensor<2xui32>
    %954 = "mhlo.broadcast_in_dim"(%873) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %955 = mhlo.add %949, %954 : tensor<2xui32>
    %956 = "mhlo.broadcast_in_dim"(%875) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %957 = mhlo.add %953, %956 : tensor<2xui32>
    %958 = mhlo.add %957, %16 : tensor<2xui32>
    %959 = mhlo.add %955, %958 : tensor<2xui32>
    %960 = mhlo.shift_left %958, %8 : tensor<2xui32>
    %961 = mhlo.shift_right_logical %958, %9 : tensor<2xui32>
    %962 = mhlo.or %960, %961 : tensor<2xui32>
    %963 = mhlo.xor %959, %962 : tensor<2xui32>
    %964 = mhlo.add %959, %963 : tensor<2xui32>
    %965 = mhlo.shift_left %963, %17 : tensor<2xui32>
    %966 = mhlo.shift_right_logical %963, %16 : tensor<2xui32>
    %967 = mhlo.or %965, %966 : tensor<2xui32>
    %968 = mhlo.xor %964, %967 : tensor<2xui32>
    %969 = mhlo.add %964, %968 : tensor<2xui32>
    %970 = mhlo.shift_left %968, %15 : tensor<2xui32>
    %971 = mhlo.shift_right_logical %968, %15 : tensor<2xui32>
    %972 = mhlo.or %970, %971 : tensor<2xui32>
    %973 = mhlo.xor %969, %972 : tensor<2xui32>
    %974 = mhlo.add %969, %973 : tensor<2xui32>
    %975 = mhlo.shift_left %973, %14 : tensor<2xui32>
    %976 = mhlo.shift_right_logical %973, %13 : tensor<2xui32>
    %977 = mhlo.or %975, %976 : tensor<2xui32>
    %978 = mhlo.xor %974, %977 : tensor<2xui32>
    %979 = "mhlo.broadcast_in_dim"(%875) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %980 = mhlo.add %974, %979 : tensor<2xui32>
    %981 = "mhlo.broadcast_in_dim"(%879) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %982 = mhlo.add %978, %981 : tensor<2xui32>
    %983 = mhlo.add %982, %12 : tensor<2xui32>
    %984 = mhlo.add %980, %983 : tensor<2xui32>
    %985 = mhlo.shift_left %983, %11 : tensor<2xui32>
    %986 = mhlo.shift_right_logical %983, %10 : tensor<2xui32>
    %987 = mhlo.or %985, %986 : tensor<2xui32>
    %988 = mhlo.xor %984, %987 : tensor<2xui32>
    %989 = mhlo.add %984, %988 : tensor<2xui32>
    %990 = mhlo.shift_left %988, %9 : tensor<2xui32>
    %991 = mhlo.shift_right_logical %988, %8 : tensor<2xui32>
    %992 = mhlo.or %990, %991 : tensor<2xui32>
    %993 = mhlo.xor %989, %992 : tensor<2xui32>
    %994 = mhlo.add %989, %993 : tensor<2xui32>
    %995 = mhlo.shift_left %993, %6 : tensor<2xui32>
    %996 = mhlo.shift_right_logical %993, %7 : tensor<2xui32>
    %997 = mhlo.or %995, %996 : tensor<2xui32>
    %998 = mhlo.xor %994, %997 : tensor<2xui32>
    %999 = mhlo.add %994, %998 : tensor<2xui32>
    %1000 = mhlo.shift_left %998, %7 : tensor<2xui32>
    %1001 = mhlo.shift_right_logical %998, %6 : tensor<2xui32>
    %1002 = mhlo.or %1000, %1001 : tensor<2xui32>
    %1003 = mhlo.xor %999, %1002 : tensor<2xui32>
    %1004 = "mhlo.broadcast_in_dim"(%879) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1005 = mhlo.add %999, %1004 : tensor<2xui32>
    %1006 = "mhlo.broadcast_in_dim"(%873) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1007 = mhlo.add %1003, %1006 : tensor<2xui32>
    %1008 = mhlo.add %1007, %5 : tensor<2xui32>
    %1009 = "mhlo.concatenate"(%1005, %1008) {dimension = 0 : i64} : (tensor<2xui32>, tensor<2xui32>) -> tensor<4xui32>
    %1010 = "mhlo.reshape"(%1009) : (tensor<4xui32>) -> tensor<2x2xui32>
    %1011 = "mhlo.slice"(%1010) {limit_indices = dense<[1, 2]> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<2x2xui32>) -> tensor<1x2xui32>
    %1012 = "mhlo.reshape"(%1011) : (tensor<1x2xui32>) -> tensor<2xui32>
    %1013 = "mhlo.slice"(%1010) {limit_indices = dense<2> : tensor<2xi64>, start_indices = dense<[1, 0]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<2x2xui32>) -> tensor<1x2xui32>
    %1014 = "mhlo.reshape"(%1013) : (tensor<1x2xui32>) -> tensor<2xui32>
    %1015 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<4xui32>
    %1016 = "mhlo.slice"(%1014) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xui32>) -> tensor<1xui32>
    %1017 = "mhlo.reshape"(%1016) : (tensor<1xui32>) -> tensor<ui32>
    %1018 = "mhlo.slice"(%1014) {limit_indices = dense<2> : tensor<1xi64>, start_indices = dense<1> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xui32>) -> tensor<1xui32>
    %1019 = "mhlo.reshape"(%1018) : (tensor<1xui32>) -> tensor<ui32>
    %1020 = "mhlo.slice"(%1015) {limit_indices = dense<2> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<4xui32>) -> tensor<2xui32>
    %1021 = "mhlo.slice"(%1015) {limit_indices = dense<4> : tensor<1xi64>, start_indices = dense<2> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<4xui32>) -> tensor<2xui32>
    %1022 = mhlo.xor %1017, %1019 : tensor<ui32>
    %1023 = mhlo.xor %1022, %20 : tensor<ui32>
    %1024 = "mhlo.broadcast_in_dim"(%1017) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1025 = mhlo.add %1020, %1024 : tensor<2xui32>
    %1026 = "mhlo.broadcast_in_dim"(%1019) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1027 = mhlo.add %1021, %1026 : tensor<2xui32>
    %1028 = mhlo.add %1025, %1027 : tensor<2xui32>
    %1029 = mhlo.shift_left %1027, %11 : tensor<2xui32>
    %1030 = mhlo.shift_right_logical %1027, %10 : tensor<2xui32>
    %1031 = mhlo.or %1029, %1030 : tensor<2xui32>
    %1032 = mhlo.xor %1028, %1031 : tensor<2xui32>
    %1033 = mhlo.add %1028, %1032 : tensor<2xui32>
    %1034 = mhlo.shift_left %1032, %9 : tensor<2xui32>
    %1035 = mhlo.shift_right_logical %1032, %8 : tensor<2xui32>
    %1036 = mhlo.or %1034, %1035 : tensor<2xui32>
    %1037 = mhlo.xor %1033, %1036 : tensor<2xui32>
    %1038 = mhlo.add %1033, %1037 : tensor<2xui32>
    %1039 = mhlo.shift_left %1037, %6 : tensor<2xui32>
    %1040 = mhlo.shift_right_logical %1037, %7 : tensor<2xui32>
    %1041 = mhlo.or %1039, %1040 : tensor<2xui32>
    %1042 = mhlo.xor %1038, %1041 : tensor<2xui32>
    %1043 = mhlo.add %1038, %1042 : tensor<2xui32>
    %1044 = mhlo.shift_left %1042, %7 : tensor<2xui32>
    %1045 = mhlo.shift_right_logical %1042, %6 : tensor<2xui32>
    %1046 = mhlo.or %1044, %1045 : tensor<2xui32>
    %1047 = mhlo.xor %1043, %1046 : tensor<2xui32>
    %1048 = "mhlo.broadcast_in_dim"(%1019) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1049 = mhlo.add %1043, %1048 : tensor<2xui32>
    %1050 = "mhlo.broadcast_in_dim"(%1023) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1051 = mhlo.add %1047, %1050 : tensor<2xui32>
    %1052 = mhlo.add %1051, %19 : tensor<2xui32>
    %1053 = mhlo.add %1049, %1052 : tensor<2xui32>
    %1054 = mhlo.shift_left %1052, %8 : tensor<2xui32>
    %1055 = mhlo.shift_right_logical %1052, %9 : tensor<2xui32>
    %1056 = mhlo.or %1054, %1055 : tensor<2xui32>
    %1057 = mhlo.xor %1053, %1056 : tensor<2xui32>
    %1058 = mhlo.add %1053, %1057 : tensor<2xui32>
    %1059 = mhlo.shift_left %1057, %17 : tensor<2xui32>
    %1060 = mhlo.shift_right_logical %1057, %16 : tensor<2xui32>
    %1061 = mhlo.or %1059, %1060 : tensor<2xui32>
    %1062 = mhlo.xor %1058, %1061 : tensor<2xui32>
    %1063 = mhlo.add %1058, %1062 : tensor<2xui32>
    %1064 = mhlo.shift_left %1062, %15 : tensor<2xui32>
    %1065 = mhlo.shift_right_logical %1062, %15 : tensor<2xui32>
    %1066 = mhlo.or %1064, %1065 : tensor<2xui32>
    %1067 = mhlo.xor %1063, %1066 : tensor<2xui32>
    %1068 = mhlo.add %1063, %1067 : tensor<2xui32>
    %1069 = mhlo.shift_left %1067, %14 : tensor<2xui32>
    %1070 = mhlo.shift_right_logical %1067, %13 : tensor<2xui32>
    %1071 = mhlo.or %1069, %1070 : tensor<2xui32>
    %1072 = mhlo.xor %1068, %1071 : tensor<2xui32>
    %1073 = "mhlo.broadcast_in_dim"(%1023) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1074 = mhlo.add %1068, %1073 : tensor<2xui32>
    %1075 = "mhlo.broadcast_in_dim"(%1017) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1076 = mhlo.add %1072, %1075 : tensor<2xui32>
    %1077 = mhlo.add %1076, %18 : tensor<2xui32>
    %1078 = mhlo.add %1074, %1077 : tensor<2xui32>
    %1079 = mhlo.shift_left %1077, %11 : tensor<2xui32>
    %1080 = mhlo.shift_right_logical %1077, %10 : tensor<2xui32>
    %1081 = mhlo.or %1079, %1080 : tensor<2xui32>
    %1082 = mhlo.xor %1078, %1081 : tensor<2xui32>
    %1083 = mhlo.add %1078, %1082 : tensor<2xui32>
    %1084 = mhlo.shift_left %1082, %9 : tensor<2xui32>
    %1085 = mhlo.shift_right_logical %1082, %8 : tensor<2xui32>
    %1086 = mhlo.or %1084, %1085 : tensor<2xui32>
    %1087 = mhlo.xor %1083, %1086 : tensor<2xui32>
    %1088 = mhlo.add %1083, %1087 : tensor<2xui32>
    %1089 = mhlo.shift_left %1087, %6 : tensor<2xui32>
    %1090 = mhlo.shift_right_logical %1087, %7 : tensor<2xui32>
    %1091 = mhlo.or %1089, %1090 : tensor<2xui32>
    %1092 = mhlo.xor %1088, %1091 : tensor<2xui32>
    %1093 = mhlo.add %1088, %1092 : tensor<2xui32>
    %1094 = mhlo.shift_left %1092, %7 : tensor<2xui32>
    %1095 = mhlo.shift_right_logical %1092, %6 : tensor<2xui32>
    %1096 = mhlo.or %1094, %1095 : tensor<2xui32>
    %1097 = mhlo.xor %1093, %1096 : tensor<2xui32>
    %1098 = "mhlo.broadcast_in_dim"(%1017) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1099 = mhlo.add %1093, %1098 : tensor<2xui32>
    %1100 = "mhlo.broadcast_in_dim"(%1019) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1101 = mhlo.add %1097, %1100 : tensor<2xui32>
    %1102 = mhlo.add %1101, %16 : tensor<2xui32>
    %1103 = mhlo.add %1099, %1102 : tensor<2xui32>
    %1104 = mhlo.shift_left %1102, %8 : tensor<2xui32>
    %1105 = mhlo.shift_right_logical %1102, %9 : tensor<2xui32>
    %1106 = mhlo.or %1104, %1105 : tensor<2xui32>
    %1107 = mhlo.xor %1103, %1106 : tensor<2xui32>
    %1108 = mhlo.add %1103, %1107 : tensor<2xui32>
    %1109 = mhlo.shift_left %1107, %17 : tensor<2xui32>
    %1110 = mhlo.shift_right_logical %1107, %16 : tensor<2xui32>
    %1111 = mhlo.or %1109, %1110 : tensor<2xui32>
    %1112 = mhlo.xor %1108, %1111 : tensor<2xui32>
    %1113 = mhlo.add %1108, %1112 : tensor<2xui32>
    %1114 = mhlo.shift_left %1112, %15 : tensor<2xui32>
    %1115 = mhlo.shift_right_logical %1112, %15 : tensor<2xui32>
    %1116 = mhlo.or %1114, %1115 : tensor<2xui32>
    %1117 = mhlo.xor %1113, %1116 : tensor<2xui32>
    %1118 = mhlo.add %1113, %1117 : tensor<2xui32>
    %1119 = mhlo.shift_left %1117, %14 : tensor<2xui32>
    %1120 = mhlo.shift_right_logical %1117, %13 : tensor<2xui32>
    %1121 = mhlo.or %1119, %1120 : tensor<2xui32>
    %1122 = mhlo.xor %1118, %1121 : tensor<2xui32>
    %1123 = "mhlo.broadcast_in_dim"(%1019) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1124 = mhlo.add %1118, %1123 : tensor<2xui32>
    %1125 = "mhlo.broadcast_in_dim"(%1023) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1126 = mhlo.add %1122, %1125 : tensor<2xui32>
    %1127 = mhlo.add %1126, %12 : tensor<2xui32>
    %1128 = mhlo.add %1124, %1127 : tensor<2xui32>
    %1129 = mhlo.shift_left %1127, %11 : tensor<2xui32>
    %1130 = mhlo.shift_right_logical %1127, %10 : tensor<2xui32>
    %1131 = mhlo.or %1129, %1130 : tensor<2xui32>
    %1132 = mhlo.xor %1128, %1131 : tensor<2xui32>
    %1133 = mhlo.add %1128, %1132 : tensor<2xui32>
    %1134 = mhlo.shift_left %1132, %9 : tensor<2xui32>
    %1135 = mhlo.shift_right_logical %1132, %8 : tensor<2xui32>
    %1136 = mhlo.or %1134, %1135 : tensor<2xui32>
    %1137 = mhlo.xor %1133, %1136 : tensor<2xui32>
    %1138 = mhlo.add %1133, %1137 : tensor<2xui32>
    %1139 = mhlo.shift_left %1137, %6 : tensor<2xui32>
    %1140 = mhlo.shift_right_logical %1137, %7 : tensor<2xui32>
    %1141 = mhlo.or %1139, %1140 : tensor<2xui32>
    %1142 = mhlo.xor %1138, %1141 : tensor<2xui32>
    %1143 = mhlo.add %1138, %1142 : tensor<2xui32>
    %1144 = mhlo.shift_left %1142, %7 : tensor<2xui32>
    %1145 = mhlo.shift_right_logical %1142, %6 : tensor<2xui32>
    %1146 = mhlo.or %1144, %1145 : tensor<2xui32>
    %1147 = mhlo.xor %1143, %1146 : tensor<2xui32>
    %1148 = "mhlo.broadcast_in_dim"(%1023) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1149 = mhlo.add %1143, %1148 : tensor<2xui32>
    %1150 = "mhlo.broadcast_in_dim"(%1017) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1151 = mhlo.add %1147, %1150 : tensor<2xui32>
    %1152 = mhlo.add %1151, %5 : tensor<2xui32>
    %1153 = "mhlo.concatenate"(%1149, %1152) {dimension = 0 : i64} : (tensor<2xui32>, tensor<2xui32>) -> tensor<4xui32>
    %1154 = "mhlo.reshape"(%1153) : (tensor<4xui32>) -> tensor<2x2xui32>
    %1155 = "mhlo.slice"(%1154) {limit_indices = dense<[1, 2]> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<2x2xui32>) -> tensor<1x2xui32>
    %1156 = "mhlo.reshape"(%1155) : (tensor<1x2xui32>) -> tensor<2xui32>
    %1157 = "mhlo.slice"(%1154) {limit_indices = dense<2> : tensor<2xi64>, start_indices = dense<[1, 0]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<2x2xui32>) -> tensor<1x2xui32>
    %1158 = "mhlo.reshape"(%1157) : (tensor<1x2xui32>) -> tensor<2xui32>
    %1159 = call @main_4(%68) : (tensor<f32>) -> tensor<f32>
    %1160 = call @main_5(%23) : (tensor<f32>) -> tensor<f32>
    %1161 = "mhlo.broadcast_in_dim"(%1159) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1x1xf32>
    %1162 = "mhlo.broadcast_in_dim"(%1160) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1x1xf32>
    %1163 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<1048576xui32>
    %1164 = "mhlo.slice"(%1156) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xui32>) -> tensor<1xui32>
    %1165 = "mhlo.reshape"(%1164) : (tensor<1xui32>) -> tensor<ui32>
    %1166 = "mhlo.slice"(%1156) {limit_indices = dense<2> : tensor<1xi64>, start_indices = dense<1> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xui32>) -> tensor<1xui32>
    %1167 = "mhlo.reshape"(%1166) : (tensor<1xui32>) -> tensor<ui32>
    %1168 = "mhlo.slice"(%1163) {limit_indices = dense<524288> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1048576xui32>) -> tensor<524288xui32>
    %1169 = "mhlo.slice"(%1163) {limit_indices = dense<1048576> : tensor<1xi64>, start_indices = dense<524288> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1048576xui32>) -> tensor<524288xui32>
    %1170 = mhlo.xor %1165, %1167 : tensor<ui32>
    %1171 = mhlo.xor %1170, %20 : tensor<ui32>
    %1172 = "mhlo.broadcast_in_dim"(%1165) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<524288xui32>
    %1173 = mhlo.add %1168, %1172 : tensor<524288xui32>
    %1174 = "mhlo.broadcast_in_dim"(%1167) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<524288xui32>
    %1175 = mhlo.add %1169, %1174 : tensor<524288xui32>
    %1176 = mhlo.add %1173, %1175 : tensor<524288xui32>
    %1177 = mhlo.shift_left %1175, %102 : tensor<524288xui32>
    %1178 = mhlo.shift_right_logical %1175, %101 : tensor<524288xui32>
    %1179 = mhlo.or %1177, %1178 : tensor<524288xui32>
    %1180 = mhlo.xor %1176, %1179 : tensor<524288xui32>
    %1181 = mhlo.add %1176, %1180 : tensor<524288xui32>
    %1182 = mhlo.shift_left %1180, %100 : tensor<524288xui32>
    %1183 = mhlo.shift_right_logical %1180, %99 : tensor<524288xui32>
    %1184 = mhlo.or %1182, %1183 : tensor<524288xui32>
    %1185 = mhlo.xor %1181, %1184 : tensor<524288xui32>
    %1186 = mhlo.add %1181, %1185 : tensor<524288xui32>
    %1187 = mhlo.shift_left %1185, %97 : tensor<524288xui32>
    %1188 = mhlo.shift_right_logical %1185, %98 : tensor<524288xui32>
    %1189 = mhlo.or %1187, %1188 : tensor<524288xui32>
    %1190 = mhlo.xor %1186, %1189 : tensor<524288xui32>
    %1191 = mhlo.add %1186, %1190 : tensor<524288xui32>
    %1192 = mhlo.shift_left %1190, %98 : tensor<524288xui32>
    %1193 = mhlo.shift_right_logical %1190, %97 : tensor<524288xui32>
    %1194 = mhlo.or %1192, %1193 : tensor<524288xui32>
    %1195 = mhlo.xor %1191, %1194 : tensor<524288xui32>
    %1196 = "mhlo.broadcast_in_dim"(%1167) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<524288xui32>
    %1197 = mhlo.add %1191, %1196 : tensor<524288xui32>
    %1198 = "mhlo.broadcast_in_dim"(%1171) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<524288xui32>
    %1199 = mhlo.add %1195, %1198 : tensor<524288xui32>
    %1200 = mhlo.add %1199, %110 : tensor<524288xui32>
    %1201 = mhlo.add %1197, %1200 : tensor<524288xui32>
    %1202 = mhlo.shift_left %1200, %99 : tensor<524288xui32>
    %1203 = mhlo.shift_right_logical %1200, %100 : tensor<524288xui32>
    %1204 = mhlo.or %1202, %1203 : tensor<524288xui32>
    %1205 = mhlo.xor %1201, %1204 : tensor<524288xui32>
    %1206 = mhlo.add %1201, %1205 : tensor<524288xui32>
    %1207 = mhlo.shift_left %1205, %108 : tensor<524288xui32>
    %1208 = mhlo.shift_right_logical %1205, %107 : tensor<524288xui32>
    %1209 = mhlo.or %1207, %1208 : tensor<524288xui32>
    %1210 = mhlo.xor %1206, %1209 : tensor<524288xui32>
    %1211 = mhlo.add %1206, %1210 : tensor<524288xui32>
    %1212 = mhlo.shift_left %1210, %106 : tensor<524288xui32>
    %1213 = mhlo.shift_right_logical %1210, %106 : tensor<524288xui32>
    %1214 = mhlo.or %1212, %1213 : tensor<524288xui32>
    %1215 = mhlo.xor %1211, %1214 : tensor<524288xui32>
    %1216 = mhlo.add %1211, %1215 : tensor<524288xui32>
    %1217 = mhlo.shift_left %1215, %105 : tensor<524288xui32>
    %1218 = mhlo.shift_right_logical %1215, %104 : tensor<524288xui32>
    %1219 = mhlo.or %1217, %1218 : tensor<524288xui32>
    %1220 = mhlo.xor %1216, %1219 : tensor<524288xui32>
    %1221 = "mhlo.broadcast_in_dim"(%1171) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<524288xui32>
    %1222 = mhlo.add %1216, %1221 : tensor<524288xui32>
    %1223 = "mhlo.broadcast_in_dim"(%1165) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<524288xui32>
    %1224 = mhlo.add %1220, %1223 : tensor<524288xui32>
    %1225 = mhlo.add %1224, %109 : tensor<524288xui32>
    %1226 = mhlo.add %1222, %1225 : tensor<524288xui32>
    %1227 = mhlo.shift_left %1225, %102 : tensor<524288xui32>
    %1228 = mhlo.shift_right_logical %1225, %101 : tensor<524288xui32>
    %1229 = mhlo.or %1227, %1228 : tensor<524288xui32>
    %1230 = mhlo.xor %1226, %1229 : tensor<524288xui32>
    %1231 = mhlo.add %1226, %1230 : tensor<524288xui32>
    %1232 = mhlo.shift_left %1230, %100 : tensor<524288xui32>
    %1233 = mhlo.shift_right_logical %1230, %99 : tensor<524288xui32>
    %1234 = mhlo.or %1232, %1233 : tensor<524288xui32>
    %1235 = mhlo.xor %1231, %1234 : tensor<524288xui32>
    %1236 = mhlo.add %1231, %1235 : tensor<524288xui32>
    %1237 = mhlo.shift_left %1235, %97 : tensor<524288xui32>
    %1238 = mhlo.shift_right_logical %1235, %98 : tensor<524288xui32>
    %1239 = mhlo.or %1237, %1238 : tensor<524288xui32>
    %1240 = mhlo.xor %1236, %1239 : tensor<524288xui32>
    %1241 = mhlo.add %1236, %1240 : tensor<524288xui32>
    %1242 = mhlo.shift_left %1240, %98 : tensor<524288xui32>
    %1243 = mhlo.shift_right_logical %1240, %97 : tensor<524288xui32>
    %1244 = mhlo.or %1242, %1243 : tensor<524288xui32>
    %1245 = mhlo.xor %1241, %1244 : tensor<524288xui32>
    %1246 = "mhlo.broadcast_in_dim"(%1165) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<524288xui32>
    %1247 = mhlo.add %1241, %1246 : tensor<524288xui32>
    %1248 = "mhlo.broadcast_in_dim"(%1167) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<524288xui32>
    %1249 = mhlo.add %1245, %1248 : tensor<524288xui32>
    %1250 = mhlo.add %1249, %107 : tensor<524288xui32>
    %1251 = mhlo.add %1247, %1250 : tensor<524288xui32>
    %1252 = mhlo.shift_left %1250, %99 : tensor<524288xui32>
    %1253 = mhlo.shift_right_logical %1250, %100 : tensor<524288xui32>
    %1254 = mhlo.or %1252, %1253 : tensor<524288xui32>
    %1255 = mhlo.xor %1251, %1254 : tensor<524288xui32>
    %1256 = mhlo.add %1251, %1255 : tensor<524288xui32>
    %1257 = mhlo.shift_left %1255, %108 : tensor<524288xui32>
    %1258 = mhlo.shift_right_logical %1255, %107 : tensor<524288xui32>
    %1259 = mhlo.or %1257, %1258 : tensor<524288xui32>
    %1260 = mhlo.xor %1256, %1259 : tensor<524288xui32>
    %1261 = mhlo.add %1256, %1260 : tensor<524288xui32>
    %1262 = mhlo.shift_left %1260, %106 : tensor<524288xui32>
    %1263 = mhlo.shift_right_logical %1260, %106 : tensor<524288xui32>
    %1264 = mhlo.or %1262, %1263 : tensor<524288xui32>
    %1265 = mhlo.xor %1261, %1264 : tensor<524288xui32>
    %1266 = mhlo.add %1261, %1265 : tensor<524288xui32>
    %1267 = mhlo.shift_left %1265, %105 : tensor<524288xui32>
    %1268 = mhlo.shift_right_logical %1265, %104 : tensor<524288xui32>
    %1269 = mhlo.or %1267, %1268 : tensor<524288xui32>
    %1270 = mhlo.xor %1266, %1269 : tensor<524288xui32>
    %1271 = "mhlo.broadcast_in_dim"(%1167) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<524288xui32>
    %1272 = mhlo.add %1266, %1271 : tensor<524288xui32>
    %1273 = "mhlo.broadcast_in_dim"(%1171) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<524288xui32>
    %1274 = mhlo.add %1270, %1273 : tensor<524288xui32>
    %1275 = mhlo.add %1274, %103 : tensor<524288xui32>
    %1276 = mhlo.add %1272, %1275 : tensor<524288xui32>
    %1277 = mhlo.shift_left %1275, %102 : tensor<524288xui32>
    %1278 = mhlo.shift_right_logical %1275, %101 : tensor<524288xui32>
    %1279 = mhlo.or %1277, %1278 : tensor<524288xui32>
    %1280 = mhlo.xor %1276, %1279 : tensor<524288xui32>
    %1281 = mhlo.add %1276, %1280 : tensor<524288xui32>
    %1282 = mhlo.shift_left %1280, %100 : tensor<524288xui32>
    %1283 = mhlo.shift_right_logical %1280, %99 : tensor<524288xui32>
    %1284 = mhlo.or %1282, %1283 : tensor<524288xui32>
    %1285 = mhlo.xor %1281, %1284 : tensor<524288xui32>
    %1286 = mhlo.add %1281, %1285 : tensor<524288xui32>
    %1287 = mhlo.shift_left %1285, %97 : tensor<524288xui32>
    %1288 = mhlo.shift_right_logical %1285, %98 : tensor<524288xui32>
    %1289 = mhlo.or %1287, %1288 : tensor<524288xui32>
    %1290 = mhlo.xor %1286, %1289 : tensor<524288xui32>
    %1291 = mhlo.add %1286, %1290 : tensor<524288xui32>
    %1292 = mhlo.shift_left %1290, %98 : tensor<524288xui32>
    %1293 = mhlo.shift_right_logical %1290, %97 : tensor<524288xui32>
    %1294 = mhlo.or %1292, %1293 : tensor<524288xui32>
    %1295 = mhlo.xor %1291, %1294 : tensor<524288xui32>
    %1296 = "mhlo.broadcast_in_dim"(%1171) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<524288xui32>
    %1297 = mhlo.add %1291, %1296 : tensor<524288xui32>
    %1298 = "mhlo.broadcast_in_dim"(%1165) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<524288xui32>
    %1299 = mhlo.add %1295, %1298 : tensor<524288xui32>
    %1300 = mhlo.add %1299, %96 : tensor<524288xui32>
    %1301 = "mhlo.concatenate"(%1297, %1300) {dimension = 0 : i64} : (tensor<524288xui32>, tensor<524288xui32>) -> tensor<1048576xui32>
    %1302 = "mhlo.reshape"(%1301) : (tensor<1048576xui32>) -> tensor<1024x1024xui32>
    %1303 = mhlo.shift_right_logical %1302, %95 : tensor<1024x1024xui32>
    %1304 = mhlo.or %1303, %94 : tensor<1024x1024xui32>
    %1305 = "mhlo.bitcast_convert"(%1304) : (tensor<1024x1024xui32>) -> tensor<1024x1024xf32>
    %1306 = mhlo.subtract %1305, %93 : tensor<1024x1024xf32>
    %1307 = mhlo.subtract %1162, %1161 : tensor<1x1xf32>
    %1308 = "mhlo.broadcast_in_dim"(%1307) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor<1x1xf32>) -> tensor<1024x1024xf32>
    %1309 = mhlo.multiply %1306, %1308 : tensor<1024x1024xf32>
    %1310 = "mhlo.broadcast_in_dim"(%1161) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor<1x1xf32>) -> tensor<1024x1024xf32>
    %1311 = mhlo.add %1309, %1310 : tensor<1024x1024xf32>
    %1312 = "mhlo.broadcast_in_dim"(%1161) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor<1x1xf32>) -> tensor<1024x1024xf32>
    %1313 = mhlo.maximum %1312, %1311 : tensor<1024x1024xf32>
    %1314 = call @main_6(%1313) : (tensor<1024x1024xf32>) -> tensor<1024x1024xf32>
    %1315 = mhlo.multiply %1314, %92 : tensor<1024x1024xf32>
    %1316 = chlo.next_after %48, %47 : tensor<f32>, tensor<f32> -> tensor<f32>
    %1317 = chlo.next_after %46, %45 : tensor<f32>, tensor<f32> -> tensor<f32>
    %1318 = call @jit_clip_7(%1315, %1316, %1317) : (tensor<1024x1024xf32>, tensor<f32>, tensor<f32>) -> tensor<1024x1024xf32>
    %1319 = mhlo.multiply %1318, %91 : tensor<1024x1024xf32>
    %1320 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<1024xui32>
    %1321 = "mhlo.slice"(%1158) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xui32>) -> tensor<1xui32>
    %1322 = "mhlo.reshape"(%1321) : (tensor<1xui32>) -> tensor<ui32>
    %1323 = "mhlo.slice"(%1158) {limit_indices = dense<2> : tensor<1xi64>, start_indices = dense<1> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xui32>) -> tensor<1xui32>
    %1324 = "mhlo.reshape"(%1323) : (tensor<1xui32>) -> tensor<ui32>
    %1325 = "mhlo.slice"(%1320) {limit_indices = dense<512> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1024xui32>) -> tensor<512xui32>
    %1326 = "mhlo.slice"(%1320) {limit_indices = dense<1024> : tensor<1xi64>, start_indices = dense<512> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<1024xui32>) -> tensor<512xui32>
    %1327 = mhlo.xor %1322, %1324 : tensor<ui32>
    %1328 = mhlo.xor %1327, %20 : tensor<ui32>
    %1329 = "mhlo.broadcast_in_dim"(%1322) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<512xui32>
    %1330 = mhlo.add %1325, %1329 : tensor<512xui32>
    %1331 = "mhlo.broadcast_in_dim"(%1324) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<512xui32>
    %1332 = mhlo.add %1326, %1331 : tensor<512xui32>
    %1333 = mhlo.add %1330, %1332 : tensor<512xui32>
    %1334 = mhlo.shift_left %1332, %82 : tensor<512xui32>
    %1335 = mhlo.shift_right_logical %1332, %81 : tensor<512xui32>
    %1336 = mhlo.or %1334, %1335 : tensor<512xui32>
    %1337 = mhlo.xor %1333, %1336 : tensor<512xui32>
    %1338 = mhlo.add %1333, %1337 : tensor<512xui32>
    %1339 = mhlo.shift_left %1337, %80 : tensor<512xui32>
    %1340 = mhlo.shift_right_logical %1337, %79 : tensor<512xui32>
    %1341 = mhlo.or %1339, %1340 : tensor<512xui32>
    %1342 = mhlo.xor %1338, %1341 : tensor<512xui32>
    %1343 = mhlo.add %1338, %1342 : tensor<512xui32>
    %1344 = mhlo.shift_left %1342, %77 : tensor<512xui32>
    %1345 = mhlo.shift_right_logical %1342, %78 : tensor<512xui32>
    %1346 = mhlo.or %1344, %1345 : tensor<512xui32>
    %1347 = mhlo.xor %1343, %1346 : tensor<512xui32>
    %1348 = mhlo.add %1343, %1347 : tensor<512xui32>
    %1349 = mhlo.shift_left %1347, %78 : tensor<512xui32>
    %1350 = mhlo.shift_right_logical %1347, %77 : tensor<512xui32>
    %1351 = mhlo.or %1349, %1350 : tensor<512xui32>
    %1352 = mhlo.xor %1348, %1351 : tensor<512xui32>
    %1353 = "mhlo.broadcast_in_dim"(%1324) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<512xui32>
    %1354 = mhlo.add %1348, %1353 : tensor<512xui32>
    %1355 = "mhlo.broadcast_in_dim"(%1328) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<512xui32>
    %1356 = mhlo.add %1352, %1355 : tensor<512xui32>
    %1357 = mhlo.add %1356, %90 : tensor<512xui32>
    %1358 = mhlo.add %1354, %1357 : tensor<512xui32>
    %1359 = mhlo.shift_left %1357, %79 : tensor<512xui32>
    %1360 = mhlo.shift_right_logical %1357, %80 : tensor<512xui32>
    %1361 = mhlo.or %1359, %1360 : tensor<512xui32>
    %1362 = mhlo.xor %1358, %1361 : tensor<512xui32>
    %1363 = mhlo.add %1358, %1362 : tensor<512xui32>
    %1364 = mhlo.shift_left %1362, %88 : tensor<512xui32>
    %1365 = mhlo.shift_right_logical %1362, %87 : tensor<512xui32>
    %1366 = mhlo.or %1364, %1365 : tensor<512xui32>
    %1367 = mhlo.xor %1363, %1366 : tensor<512xui32>
    %1368 = mhlo.add %1363, %1367 : tensor<512xui32>
    %1369 = mhlo.shift_left %1367, %86 : tensor<512xui32>
    %1370 = mhlo.shift_right_logical %1367, %86 : tensor<512xui32>
    %1371 = mhlo.or %1369, %1370 : tensor<512xui32>
    %1372 = mhlo.xor %1368, %1371 : tensor<512xui32>
    %1373 = mhlo.add %1368, %1372 : tensor<512xui32>
    %1374 = mhlo.shift_left %1372, %85 : tensor<512xui32>
    %1375 = mhlo.shift_right_logical %1372, %84 : tensor<512xui32>
    %1376 = mhlo.or %1374, %1375 : tensor<512xui32>
    %1377 = mhlo.xor %1373, %1376 : tensor<512xui32>
    %1378 = "mhlo.broadcast_in_dim"(%1328) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<512xui32>
    %1379 = mhlo.add %1373, %1378 : tensor<512xui32>
    %1380 = "mhlo.broadcast_in_dim"(%1322) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<512xui32>
    %1381 = mhlo.add %1377, %1380 : tensor<512xui32>
    %1382 = mhlo.add %1381, %89 : tensor<512xui32>
    %1383 = mhlo.add %1379, %1382 : tensor<512xui32>
    %1384 = mhlo.shift_left %1382, %82 : tensor<512xui32>
    %1385 = mhlo.shift_right_logical %1382, %81 : tensor<512xui32>
    %1386 = mhlo.or %1384, %1385 : tensor<512xui32>
    %1387 = mhlo.xor %1383, %1386 : tensor<512xui32>
    %1388 = mhlo.add %1383, %1387 : tensor<512xui32>
    %1389 = mhlo.shift_left %1387, %80 : tensor<512xui32>
    %1390 = mhlo.shift_right_logical %1387, %79 : tensor<512xui32>
    %1391 = mhlo.or %1389, %1390 : tensor<512xui32>
    %1392 = mhlo.xor %1388, %1391 : tensor<512xui32>
    %1393 = mhlo.add %1388, %1392 : tensor<512xui32>
    %1394 = mhlo.shift_left %1392, %77 : tensor<512xui32>
    %1395 = mhlo.shift_right_logical %1392, %78 : tensor<512xui32>
    %1396 = mhlo.or %1394, %1395 : tensor<512xui32>
    %1397 = mhlo.xor %1393, %1396 : tensor<512xui32>
    %1398 = mhlo.add %1393, %1397 : tensor<512xui32>
    %1399 = mhlo.shift_left %1397, %78 : tensor<512xui32>
    %1400 = mhlo.shift_right_logical %1397, %77 : tensor<512xui32>
    %1401 = mhlo.or %1399, %1400 : tensor<512xui32>
    %1402 = mhlo.xor %1398, %1401 : tensor<512xui32>
    %1403 = "mhlo.broadcast_in_dim"(%1322) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<512xui32>
    %1404 = mhlo.add %1398, %1403 : tensor<512xui32>
    %1405 = "mhlo.broadcast_in_dim"(%1324) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<512xui32>
    %1406 = mhlo.add %1402, %1405 : tensor<512xui32>
    %1407 = mhlo.add %1406, %87 : tensor<512xui32>
    %1408 = mhlo.add %1404, %1407 : tensor<512xui32>
    %1409 = mhlo.shift_left %1407, %79 : tensor<512xui32>
    %1410 = mhlo.shift_right_logical %1407, %80 : tensor<512xui32>
    %1411 = mhlo.or %1409, %1410 : tensor<512xui32>
    %1412 = mhlo.xor %1408, %1411 : tensor<512xui32>
    %1413 = mhlo.add %1408, %1412 : tensor<512xui32>
    %1414 = mhlo.shift_left %1412, %88 : tensor<512xui32>
    %1415 = mhlo.shift_right_logical %1412, %87 : tensor<512xui32>
    %1416 = mhlo.or %1414, %1415 : tensor<512xui32>
    %1417 = mhlo.xor %1413, %1416 : tensor<512xui32>
    %1418 = mhlo.add %1413, %1417 : tensor<512xui32>
    %1419 = mhlo.shift_left %1417, %86 : tensor<512xui32>
    %1420 = mhlo.shift_right_logical %1417, %86 : tensor<512xui32>
    %1421 = mhlo.or %1419, %1420 : tensor<512xui32>
    %1422 = mhlo.xor %1418, %1421 : tensor<512xui32>
    %1423 = mhlo.add %1418, %1422 : tensor<512xui32>
    %1424 = mhlo.shift_left %1422, %85 : tensor<512xui32>
    %1425 = mhlo.shift_right_logical %1422, %84 : tensor<512xui32>
    %1426 = mhlo.or %1424, %1425 : tensor<512xui32>
    %1427 = mhlo.xor %1423, %1426 : tensor<512xui32>
    %1428 = "mhlo.broadcast_in_dim"(%1324) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<512xui32>
    %1429 = mhlo.add %1423, %1428 : tensor<512xui32>
    %1430 = "mhlo.broadcast_in_dim"(%1328) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<512xui32>
    %1431 = mhlo.add %1427, %1430 : tensor<512xui32>
    %1432 = mhlo.add %1431, %83 : tensor<512xui32>
    %1433 = mhlo.add %1429, %1432 : tensor<512xui32>
    %1434 = mhlo.shift_left %1432, %82 : tensor<512xui32>
    %1435 = mhlo.shift_right_logical %1432, %81 : tensor<512xui32>
    %1436 = mhlo.or %1434, %1435 : tensor<512xui32>
    %1437 = mhlo.xor %1433, %1436 : tensor<512xui32>
    %1438 = mhlo.add %1433, %1437 : tensor<512xui32>
    %1439 = mhlo.shift_left %1437, %80 : tensor<512xui32>
    %1440 = mhlo.shift_right_logical %1437, %79 : tensor<512xui32>
    %1441 = mhlo.or %1439, %1440 : tensor<512xui32>
    %1442 = mhlo.xor %1438, %1441 : tensor<512xui32>
    %1443 = mhlo.add %1438, %1442 : tensor<512xui32>
    %1444 = mhlo.shift_left %1442, %77 : tensor<512xui32>
    %1445 = mhlo.shift_right_logical %1442, %78 : tensor<512xui32>
    %1446 = mhlo.or %1444, %1445 : tensor<512xui32>
    %1447 = mhlo.xor %1443, %1446 : tensor<512xui32>
    %1448 = mhlo.add %1443, %1447 : tensor<512xui32>
    %1449 = mhlo.shift_left %1447, %78 : tensor<512xui32>
    %1450 = mhlo.shift_right_logical %1447, %77 : tensor<512xui32>
    %1451 = mhlo.or %1449, %1450 : tensor<512xui32>
    %1452 = mhlo.xor %1448, %1451 : tensor<512xui32>
    %1453 = "mhlo.broadcast_in_dim"(%1328) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<512xui32>
    %1454 = mhlo.add %1448, %1453 : tensor<512xui32>
    %1455 = "mhlo.broadcast_in_dim"(%1322) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<512xui32>
    %1456 = mhlo.add %1452, %1455 : tensor<512xui32>
    %1457 = mhlo.add %1456, %76 : tensor<512xui32>
    %1458 = "mhlo.concatenate"(%1454, %1457) {dimension = 0 : i64} : (tensor<512xui32>, tensor<512xui32>) -> tensor<1024xui32>
    %1459 = mhlo.shift_right_logical %1458, %75 : tensor<1024xui32>
    %1460 = mhlo.or %1459, %74 : tensor<1024xui32>
    %1461 = "mhlo.bitcast_convert"(%1460) : (tensor<1024xui32>) -> tensor<1024xf32>
    %1462 = mhlo.subtract %1461, %73 : tensor<1024xf32>
    %1463 = mhlo.multiply %1462, %72 : tensor<1024xf32>
    %1464 = mhlo.add %1463, %71 : tensor<1024xf32>
    %1465 = mhlo.maximum %1464, %71 : tensor<1024xf32>
    %1466 = call @main_8(%1465) : (tensor<1024xf32>) -> tensor<1024xf32>
    %1467 = mhlo.multiply %1466, %70 : tensor<1024xf32>
    %1468 = mhlo.multiply %1467, %69 : tensor<1024xf32>
    %1469 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<4xui32>
    %1470 = "mhlo.slice"(%1012) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xui32>) -> tensor<1xui32>
    %1471 = "mhlo.reshape"(%1470) : (tensor<1xui32>) -> tensor<ui32>
    %1472 = "mhlo.slice"(%1012) {limit_indices = dense<2> : tensor<1xi64>, start_indices = dense<1> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xui32>) -> tensor<1xui32>
    %1473 = "mhlo.reshape"(%1472) : (tensor<1xui32>) -> tensor<ui32>
    %1474 = "mhlo.slice"(%1469) {limit_indices = dense<2> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<4xui32>) -> tensor<2xui32>
    %1475 = "mhlo.slice"(%1469) {limit_indices = dense<4> : tensor<1xi64>, start_indices = dense<2> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<4xui32>) -> tensor<2xui32>
    %1476 = mhlo.xor %1471, %1473 : tensor<ui32>
    %1477 = mhlo.xor %1476, %20 : tensor<ui32>
    %1478 = "mhlo.broadcast_in_dim"(%1471) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1479 = mhlo.add %1474, %1478 : tensor<2xui32>
    %1480 = "mhlo.broadcast_in_dim"(%1473) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1481 = mhlo.add %1475, %1480 : tensor<2xui32>
    %1482 = mhlo.add %1479, %1481 : tensor<2xui32>
    %1483 = mhlo.shift_left %1481, %11 : tensor<2xui32>
    %1484 = mhlo.shift_right_logical %1481, %10 : tensor<2xui32>
    %1485 = mhlo.or %1483, %1484 : tensor<2xui32>
    %1486 = mhlo.xor %1482, %1485 : tensor<2xui32>
    %1487 = mhlo.add %1482, %1486 : tensor<2xui32>
    %1488 = mhlo.shift_left %1486, %9 : tensor<2xui32>
    %1489 = mhlo.shift_right_logical %1486, %8 : tensor<2xui32>
    %1490 = mhlo.or %1488, %1489 : tensor<2xui32>
    %1491 = mhlo.xor %1487, %1490 : tensor<2xui32>
    %1492 = mhlo.add %1487, %1491 : tensor<2xui32>
    %1493 = mhlo.shift_left %1491, %6 : tensor<2xui32>
    %1494 = mhlo.shift_right_logical %1491, %7 : tensor<2xui32>
    %1495 = mhlo.or %1493, %1494 : tensor<2xui32>
    %1496 = mhlo.xor %1492, %1495 : tensor<2xui32>
    %1497 = mhlo.add %1492, %1496 : tensor<2xui32>
    %1498 = mhlo.shift_left %1496, %7 : tensor<2xui32>
    %1499 = mhlo.shift_right_logical %1496, %6 : tensor<2xui32>
    %1500 = mhlo.or %1498, %1499 : tensor<2xui32>
    %1501 = mhlo.xor %1497, %1500 : tensor<2xui32>
    %1502 = "mhlo.broadcast_in_dim"(%1473) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1503 = mhlo.add %1497, %1502 : tensor<2xui32>
    %1504 = "mhlo.broadcast_in_dim"(%1477) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1505 = mhlo.add %1501, %1504 : tensor<2xui32>
    %1506 = mhlo.add %1505, %19 : tensor<2xui32>
    %1507 = mhlo.add %1503, %1506 : tensor<2xui32>
    %1508 = mhlo.shift_left %1506, %8 : tensor<2xui32>
    %1509 = mhlo.shift_right_logical %1506, %9 : tensor<2xui32>
    %1510 = mhlo.or %1508, %1509 : tensor<2xui32>
    %1511 = mhlo.xor %1507, %1510 : tensor<2xui32>
    %1512 = mhlo.add %1507, %1511 : tensor<2xui32>
    %1513 = mhlo.shift_left %1511, %17 : tensor<2xui32>
    %1514 = mhlo.shift_right_logical %1511, %16 : tensor<2xui32>
    %1515 = mhlo.or %1513, %1514 : tensor<2xui32>
    %1516 = mhlo.xor %1512, %1515 : tensor<2xui32>
    %1517 = mhlo.add %1512, %1516 : tensor<2xui32>
    %1518 = mhlo.shift_left %1516, %15 : tensor<2xui32>
    %1519 = mhlo.shift_right_logical %1516, %15 : tensor<2xui32>
    %1520 = mhlo.or %1518, %1519 : tensor<2xui32>
    %1521 = mhlo.xor %1517, %1520 : tensor<2xui32>
    %1522 = mhlo.add %1517, %1521 : tensor<2xui32>
    %1523 = mhlo.shift_left %1521, %14 : tensor<2xui32>
    %1524 = mhlo.shift_right_logical %1521, %13 : tensor<2xui32>
    %1525 = mhlo.or %1523, %1524 : tensor<2xui32>
    %1526 = mhlo.xor %1522, %1525 : tensor<2xui32>
    %1527 = "mhlo.broadcast_in_dim"(%1477) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1528 = mhlo.add %1522, %1527 : tensor<2xui32>
    %1529 = "mhlo.broadcast_in_dim"(%1471) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1530 = mhlo.add %1526, %1529 : tensor<2xui32>
    %1531 = mhlo.add %1530, %18 : tensor<2xui32>
    %1532 = mhlo.add %1528, %1531 : tensor<2xui32>
    %1533 = mhlo.shift_left %1531, %11 : tensor<2xui32>
    %1534 = mhlo.shift_right_logical %1531, %10 : tensor<2xui32>
    %1535 = mhlo.or %1533, %1534 : tensor<2xui32>
    %1536 = mhlo.xor %1532, %1535 : tensor<2xui32>
    %1537 = mhlo.add %1532, %1536 : tensor<2xui32>
    %1538 = mhlo.shift_left %1536, %9 : tensor<2xui32>
    %1539 = mhlo.shift_right_logical %1536, %8 : tensor<2xui32>
    %1540 = mhlo.or %1538, %1539 : tensor<2xui32>
    %1541 = mhlo.xor %1537, %1540 : tensor<2xui32>
    %1542 = mhlo.add %1537, %1541 : tensor<2xui32>
    %1543 = mhlo.shift_left %1541, %6 : tensor<2xui32>
    %1544 = mhlo.shift_right_logical %1541, %7 : tensor<2xui32>
    %1545 = mhlo.or %1543, %1544 : tensor<2xui32>
    %1546 = mhlo.xor %1542, %1545 : tensor<2xui32>
    %1547 = mhlo.add %1542, %1546 : tensor<2xui32>
    %1548 = mhlo.shift_left %1546, %7 : tensor<2xui32>
    %1549 = mhlo.shift_right_logical %1546, %6 : tensor<2xui32>
    %1550 = mhlo.or %1548, %1549 : tensor<2xui32>
    %1551 = mhlo.xor %1547, %1550 : tensor<2xui32>
    %1552 = "mhlo.broadcast_in_dim"(%1471) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1553 = mhlo.add %1547, %1552 : tensor<2xui32>
    %1554 = "mhlo.broadcast_in_dim"(%1473) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1555 = mhlo.add %1551, %1554 : tensor<2xui32>
    %1556 = mhlo.add %1555, %16 : tensor<2xui32>
    %1557 = mhlo.add %1553, %1556 : tensor<2xui32>
    %1558 = mhlo.shift_left %1556, %8 : tensor<2xui32>
    %1559 = mhlo.shift_right_logical %1556, %9 : tensor<2xui32>
    %1560 = mhlo.or %1558, %1559 : tensor<2xui32>
    %1561 = mhlo.xor %1557, %1560 : tensor<2xui32>
    %1562 = mhlo.add %1557, %1561 : tensor<2xui32>
    %1563 = mhlo.shift_left %1561, %17 : tensor<2xui32>
    %1564 = mhlo.shift_right_logical %1561, %16 : tensor<2xui32>
    %1565 = mhlo.or %1563, %1564 : tensor<2xui32>
    %1566 = mhlo.xor %1562, %1565 : tensor<2xui32>
    %1567 = mhlo.add %1562, %1566 : tensor<2xui32>
    %1568 = mhlo.shift_left %1566, %15 : tensor<2xui32>
    %1569 = mhlo.shift_right_logical %1566, %15 : tensor<2xui32>
    %1570 = mhlo.or %1568, %1569 : tensor<2xui32>
    %1571 = mhlo.xor %1567, %1570 : tensor<2xui32>
    %1572 = mhlo.add %1567, %1571 : tensor<2xui32>
    %1573 = mhlo.shift_left %1571, %14 : tensor<2xui32>
    %1574 = mhlo.shift_right_logical %1571, %13 : tensor<2xui32>
    %1575 = mhlo.or %1573, %1574 : tensor<2xui32>
    %1576 = mhlo.xor %1572, %1575 : tensor<2xui32>
    %1577 = "mhlo.broadcast_in_dim"(%1473) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1578 = mhlo.add %1572, %1577 : tensor<2xui32>
    %1579 = "mhlo.broadcast_in_dim"(%1477) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1580 = mhlo.add %1576, %1579 : tensor<2xui32>
    %1581 = mhlo.add %1580, %12 : tensor<2xui32>
    %1582 = mhlo.add %1578, %1581 : tensor<2xui32>
    %1583 = mhlo.shift_left %1581, %11 : tensor<2xui32>
    %1584 = mhlo.shift_right_logical %1581, %10 : tensor<2xui32>
    %1585 = mhlo.or %1583, %1584 : tensor<2xui32>
    %1586 = mhlo.xor %1582, %1585 : tensor<2xui32>
    %1587 = mhlo.add %1582, %1586 : tensor<2xui32>
    %1588 = mhlo.shift_left %1586, %9 : tensor<2xui32>
    %1589 = mhlo.shift_right_logical %1586, %8 : tensor<2xui32>
    %1590 = mhlo.or %1588, %1589 : tensor<2xui32>
    %1591 = mhlo.xor %1587, %1590 : tensor<2xui32>
    %1592 = mhlo.add %1587, %1591 : tensor<2xui32>
    %1593 = mhlo.shift_left %1591, %6 : tensor<2xui32>
    %1594 = mhlo.shift_right_logical %1591, %7 : tensor<2xui32>
    %1595 = mhlo.or %1593, %1594 : tensor<2xui32>
    %1596 = mhlo.xor %1592, %1595 : tensor<2xui32>
    %1597 = mhlo.add %1592, %1596 : tensor<2xui32>
    %1598 = mhlo.shift_left %1596, %7 : tensor<2xui32>
    %1599 = mhlo.shift_right_logical %1596, %6 : tensor<2xui32>
    %1600 = mhlo.or %1598, %1599 : tensor<2xui32>
    %1601 = mhlo.xor %1597, %1600 : tensor<2xui32>
    %1602 = "mhlo.broadcast_in_dim"(%1477) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1603 = mhlo.add %1597, %1602 : tensor<2xui32>
    %1604 = "mhlo.broadcast_in_dim"(%1471) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1605 = mhlo.add %1601, %1604 : tensor<2xui32>
    %1606 = mhlo.add %1605, %5 : tensor<2xui32>
    %1607 = "mhlo.concatenate"(%1603, %1606) {dimension = 0 : i64} : (tensor<2xui32>, tensor<2xui32>) -> tensor<4xui32>
    %1608 = "mhlo.reshape"(%1607) : (tensor<4xui32>) -> tensor<2x2xui32>
    %1609 = "mhlo.slice"(%1608) {limit_indices = dense<[1, 2]> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<2x2xui32>) -> tensor<1x2xui32>
    %1610 = "mhlo.reshape"(%1609) : (tensor<1x2xui32>) -> tensor<2xui32>
    %1611 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<4xui32>
    %1612 = "mhlo.slice"(%1610) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xui32>) -> tensor<1xui32>
    %1613 = "mhlo.reshape"(%1612) : (tensor<1xui32>) -> tensor<ui32>
    %1614 = "mhlo.slice"(%1610) {limit_indices = dense<2> : tensor<1xi64>, start_indices = dense<1> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xui32>) -> tensor<1xui32>
    %1615 = "mhlo.reshape"(%1614) : (tensor<1xui32>) -> tensor<ui32>
    %1616 = "mhlo.slice"(%1611) {limit_indices = dense<2> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<4xui32>) -> tensor<2xui32>
    %1617 = "mhlo.slice"(%1611) {limit_indices = dense<4> : tensor<1xi64>, start_indices = dense<2> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<4xui32>) -> tensor<2xui32>
    %1618 = mhlo.xor %1613, %1615 : tensor<ui32>
    %1619 = mhlo.xor %1618, %20 : tensor<ui32>
    %1620 = "mhlo.broadcast_in_dim"(%1613) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1621 = mhlo.add %1616, %1620 : tensor<2xui32>
    %1622 = "mhlo.broadcast_in_dim"(%1615) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1623 = mhlo.add %1617, %1622 : tensor<2xui32>
    %1624 = mhlo.add %1621, %1623 : tensor<2xui32>
    %1625 = mhlo.shift_left %1623, %11 : tensor<2xui32>
    %1626 = mhlo.shift_right_logical %1623, %10 : tensor<2xui32>
    %1627 = mhlo.or %1625, %1626 : tensor<2xui32>
    %1628 = mhlo.xor %1624, %1627 : tensor<2xui32>
    %1629 = mhlo.add %1624, %1628 : tensor<2xui32>
    %1630 = mhlo.shift_left %1628, %9 : tensor<2xui32>
    %1631 = mhlo.shift_right_logical %1628, %8 : tensor<2xui32>
    %1632 = mhlo.or %1630, %1631 : tensor<2xui32>
    %1633 = mhlo.xor %1629, %1632 : tensor<2xui32>
    %1634 = mhlo.add %1629, %1633 : tensor<2xui32>
    %1635 = mhlo.shift_left %1633, %6 : tensor<2xui32>
    %1636 = mhlo.shift_right_logical %1633, %7 : tensor<2xui32>
    %1637 = mhlo.or %1635, %1636 : tensor<2xui32>
    %1638 = mhlo.xor %1634, %1637 : tensor<2xui32>
    %1639 = mhlo.add %1634, %1638 : tensor<2xui32>
    %1640 = mhlo.shift_left %1638, %7 : tensor<2xui32>
    %1641 = mhlo.shift_right_logical %1638, %6 : tensor<2xui32>
    %1642 = mhlo.or %1640, %1641 : tensor<2xui32>
    %1643 = mhlo.xor %1639, %1642 : tensor<2xui32>
    %1644 = "mhlo.broadcast_in_dim"(%1615) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1645 = mhlo.add %1639, %1644 : tensor<2xui32>
    %1646 = "mhlo.broadcast_in_dim"(%1619) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1647 = mhlo.add %1643, %1646 : tensor<2xui32>
    %1648 = mhlo.add %1647, %19 : tensor<2xui32>
    %1649 = mhlo.add %1645, %1648 : tensor<2xui32>
    %1650 = mhlo.shift_left %1648, %8 : tensor<2xui32>
    %1651 = mhlo.shift_right_logical %1648, %9 : tensor<2xui32>
    %1652 = mhlo.or %1650, %1651 : tensor<2xui32>
    %1653 = mhlo.xor %1649, %1652 : tensor<2xui32>
    %1654 = mhlo.add %1649, %1653 : tensor<2xui32>
    %1655 = mhlo.shift_left %1653, %17 : tensor<2xui32>
    %1656 = mhlo.shift_right_logical %1653, %16 : tensor<2xui32>
    %1657 = mhlo.or %1655, %1656 : tensor<2xui32>
    %1658 = mhlo.xor %1654, %1657 : tensor<2xui32>
    %1659 = mhlo.add %1654, %1658 : tensor<2xui32>
    %1660 = mhlo.shift_left %1658, %15 : tensor<2xui32>
    %1661 = mhlo.shift_right_logical %1658, %15 : tensor<2xui32>
    %1662 = mhlo.or %1660, %1661 : tensor<2xui32>
    %1663 = mhlo.xor %1659, %1662 : tensor<2xui32>
    %1664 = mhlo.add %1659, %1663 : tensor<2xui32>
    %1665 = mhlo.shift_left %1663, %14 : tensor<2xui32>
    %1666 = mhlo.shift_right_logical %1663, %13 : tensor<2xui32>
    %1667 = mhlo.or %1665, %1666 : tensor<2xui32>
    %1668 = mhlo.xor %1664, %1667 : tensor<2xui32>
    %1669 = "mhlo.broadcast_in_dim"(%1619) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1670 = mhlo.add %1664, %1669 : tensor<2xui32>
    %1671 = "mhlo.broadcast_in_dim"(%1613) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1672 = mhlo.add %1668, %1671 : tensor<2xui32>
    %1673 = mhlo.add %1672, %18 : tensor<2xui32>
    %1674 = mhlo.add %1670, %1673 : tensor<2xui32>
    %1675 = mhlo.shift_left %1673, %11 : tensor<2xui32>
    %1676 = mhlo.shift_right_logical %1673, %10 : tensor<2xui32>
    %1677 = mhlo.or %1675, %1676 : tensor<2xui32>
    %1678 = mhlo.xor %1674, %1677 : tensor<2xui32>
    %1679 = mhlo.add %1674, %1678 : tensor<2xui32>
    %1680 = mhlo.shift_left %1678, %9 : tensor<2xui32>
    %1681 = mhlo.shift_right_logical %1678, %8 : tensor<2xui32>
    %1682 = mhlo.or %1680, %1681 : tensor<2xui32>
    %1683 = mhlo.xor %1679, %1682 : tensor<2xui32>
    %1684 = mhlo.add %1679, %1683 : tensor<2xui32>
    %1685 = mhlo.shift_left %1683, %6 : tensor<2xui32>
    %1686 = mhlo.shift_right_logical %1683, %7 : tensor<2xui32>
    %1687 = mhlo.or %1685, %1686 : tensor<2xui32>
    %1688 = mhlo.xor %1684, %1687 : tensor<2xui32>
    %1689 = mhlo.add %1684, %1688 : tensor<2xui32>
    %1690 = mhlo.shift_left %1688, %7 : tensor<2xui32>
    %1691 = mhlo.shift_right_logical %1688, %6 : tensor<2xui32>
    %1692 = mhlo.or %1690, %1691 : tensor<2xui32>
    %1693 = mhlo.xor %1689, %1692 : tensor<2xui32>
    %1694 = "mhlo.broadcast_in_dim"(%1613) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1695 = mhlo.add %1689, %1694 : tensor<2xui32>
    %1696 = "mhlo.broadcast_in_dim"(%1615) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1697 = mhlo.add %1693, %1696 : tensor<2xui32>
    %1698 = mhlo.add %1697, %16 : tensor<2xui32>
    %1699 = mhlo.add %1695, %1698 : tensor<2xui32>
    %1700 = mhlo.shift_left %1698, %8 : tensor<2xui32>
    %1701 = mhlo.shift_right_logical %1698, %9 : tensor<2xui32>
    %1702 = mhlo.or %1700, %1701 : tensor<2xui32>
    %1703 = mhlo.xor %1699, %1702 : tensor<2xui32>
    %1704 = mhlo.add %1699, %1703 : tensor<2xui32>
    %1705 = mhlo.shift_left %1703, %17 : tensor<2xui32>
    %1706 = mhlo.shift_right_logical %1703, %16 : tensor<2xui32>
    %1707 = mhlo.or %1705, %1706 : tensor<2xui32>
    %1708 = mhlo.xor %1704, %1707 : tensor<2xui32>
    %1709 = mhlo.add %1704, %1708 : tensor<2xui32>
    %1710 = mhlo.shift_left %1708, %15 : tensor<2xui32>
    %1711 = mhlo.shift_right_logical %1708, %15 : tensor<2xui32>
    %1712 = mhlo.or %1710, %1711 : tensor<2xui32>
    %1713 = mhlo.xor %1709, %1712 : tensor<2xui32>
    %1714 = mhlo.add %1709, %1713 : tensor<2xui32>
    %1715 = mhlo.shift_left %1713, %14 : tensor<2xui32>
    %1716 = mhlo.shift_right_logical %1713, %13 : tensor<2xui32>
    %1717 = mhlo.or %1715, %1716 : tensor<2xui32>
    %1718 = mhlo.xor %1714, %1717 : tensor<2xui32>
    %1719 = "mhlo.broadcast_in_dim"(%1615) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1720 = mhlo.add %1714, %1719 : tensor<2xui32>
    %1721 = "mhlo.broadcast_in_dim"(%1619) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1722 = mhlo.add %1718, %1721 : tensor<2xui32>
    %1723 = mhlo.add %1722, %12 : tensor<2xui32>
    %1724 = mhlo.add %1720, %1723 : tensor<2xui32>
    %1725 = mhlo.shift_left %1723, %11 : tensor<2xui32>
    %1726 = mhlo.shift_right_logical %1723, %10 : tensor<2xui32>
    %1727 = mhlo.or %1725, %1726 : tensor<2xui32>
    %1728 = mhlo.xor %1724, %1727 : tensor<2xui32>
    %1729 = mhlo.add %1724, %1728 : tensor<2xui32>
    %1730 = mhlo.shift_left %1728, %9 : tensor<2xui32>
    %1731 = mhlo.shift_right_logical %1728, %8 : tensor<2xui32>
    %1732 = mhlo.or %1730, %1731 : tensor<2xui32>
    %1733 = mhlo.xor %1729, %1732 : tensor<2xui32>
    %1734 = mhlo.add %1729, %1733 : tensor<2xui32>
    %1735 = mhlo.shift_left %1733, %6 : tensor<2xui32>
    %1736 = mhlo.shift_right_logical %1733, %7 : tensor<2xui32>
    %1737 = mhlo.or %1735, %1736 : tensor<2xui32>
    %1738 = mhlo.xor %1734, %1737 : tensor<2xui32>
    %1739 = mhlo.add %1734, %1738 : tensor<2xui32>
    %1740 = mhlo.shift_left %1738, %7 : tensor<2xui32>
    %1741 = mhlo.shift_right_logical %1738, %6 : tensor<2xui32>
    %1742 = mhlo.or %1740, %1741 : tensor<2xui32>
    %1743 = mhlo.xor %1739, %1742 : tensor<2xui32>
    %1744 = "mhlo.broadcast_in_dim"(%1619) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1745 = mhlo.add %1739, %1744 : tensor<2xui32>
    %1746 = "mhlo.broadcast_in_dim"(%1613) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1747 = mhlo.add %1743, %1746 : tensor<2xui32>
    %1748 = mhlo.add %1747, %5 : tensor<2xui32>
    %1749 = "mhlo.concatenate"(%1745, %1748) {dimension = 0 : i64} : (tensor<2xui32>, tensor<2xui32>) -> tensor<4xui32>
    %1750 = "mhlo.reshape"(%1749) : (tensor<4xui32>) -> tensor<2x2xui32>
    %1751 = "mhlo.slice"(%1750) {limit_indices = dense<2> : tensor<2xi64>, start_indices = dense<[1, 0]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<2x2xui32>) -> tensor<1x2xui32>
    %1752 = "mhlo.reshape"(%1751) : (tensor<1x2xui32>) -> tensor<2xui32>
    %1753 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<4xui32>
    %1754 = "mhlo.slice"(%1752) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xui32>) -> tensor<1xui32>
    %1755 = "mhlo.reshape"(%1754) : (tensor<1xui32>) -> tensor<ui32>
    %1756 = "mhlo.slice"(%1752) {limit_indices = dense<2> : tensor<1xi64>, start_indices = dense<1> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xui32>) -> tensor<1xui32>
    %1757 = "mhlo.reshape"(%1756) : (tensor<1xui32>) -> tensor<ui32>
    %1758 = "mhlo.slice"(%1753) {limit_indices = dense<2> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<4xui32>) -> tensor<2xui32>
    %1759 = "mhlo.slice"(%1753) {limit_indices = dense<4> : tensor<1xi64>, start_indices = dense<2> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<4xui32>) -> tensor<2xui32>
    %1760 = mhlo.xor %1755, %1757 : tensor<ui32>
    %1761 = mhlo.xor %1760, %20 : tensor<ui32>
    %1762 = "mhlo.broadcast_in_dim"(%1755) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1763 = mhlo.add %1758, %1762 : tensor<2xui32>
    %1764 = "mhlo.broadcast_in_dim"(%1757) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1765 = mhlo.add %1759, %1764 : tensor<2xui32>
    %1766 = mhlo.add %1763, %1765 : tensor<2xui32>
    %1767 = mhlo.shift_left %1765, %11 : tensor<2xui32>
    %1768 = mhlo.shift_right_logical %1765, %10 : tensor<2xui32>
    %1769 = mhlo.or %1767, %1768 : tensor<2xui32>
    %1770 = mhlo.xor %1766, %1769 : tensor<2xui32>
    %1771 = mhlo.add %1766, %1770 : tensor<2xui32>
    %1772 = mhlo.shift_left %1770, %9 : tensor<2xui32>
    %1773 = mhlo.shift_right_logical %1770, %8 : tensor<2xui32>
    %1774 = mhlo.or %1772, %1773 : tensor<2xui32>
    %1775 = mhlo.xor %1771, %1774 : tensor<2xui32>
    %1776 = mhlo.add %1771, %1775 : tensor<2xui32>
    %1777 = mhlo.shift_left %1775, %6 : tensor<2xui32>
    %1778 = mhlo.shift_right_logical %1775, %7 : tensor<2xui32>
    %1779 = mhlo.or %1777, %1778 : tensor<2xui32>
    %1780 = mhlo.xor %1776, %1779 : tensor<2xui32>
    %1781 = mhlo.add %1776, %1780 : tensor<2xui32>
    %1782 = mhlo.shift_left %1780, %7 : tensor<2xui32>
    %1783 = mhlo.shift_right_logical %1780, %6 : tensor<2xui32>
    %1784 = mhlo.or %1782, %1783 : tensor<2xui32>
    %1785 = mhlo.xor %1781, %1784 : tensor<2xui32>
    %1786 = "mhlo.broadcast_in_dim"(%1757) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1787 = mhlo.add %1781, %1786 : tensor<2xui32>
    %1788 = "mhlo.broadcast_in_dim"(%1761) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1789 = mhlo.add %1785, %1788 : tensor<2xui32>
    %1790 = mhlo.add %1789, %19 : tensor<2xui32>
    %1791 = mhlo.add %1787, %1790 : tensor<2xui32>
    %1792 = mhlo.shift_left %1790, %8 : tensor<2xui32>
    %1793 = mhlo.shift_right_logical %1790, %9 : tensor<2xui32>
    %1794 = mhlo.or %1792, %1793 : tensor<2xui32>
    %1795 = mhlo.xor %1791, %1794 : tensor<2xui32>
    %1796 = mhlo.add %1791, %1795 : tensor<2xui32>
    %1797 = mhlo.shift_left %1795, %17 : tensor<2xui32>
    %1798 = mhlo.shift_right_logical %1795, %16 : tensor<2xui32>
    %1799 = mhlo.or %1797, %1798 : tensor<2xui32>
    %1800 = mhlo.xor %1796, %1799 : tensor<2xui32>
    %1801 = mhlo.add %1796, %1800 : tensor<2xui32>
    %1802 = mhlo.shift_left %1800, %15 : tensor<2xui32>
    %1803 = mhlo.shift_right_logical %1800, %15 : tensor<2xui32>
    %1804 = mhlo.or %1802, %1803 : tensor<2xui32>
    %1805 = mhlo.xor %1801, %1804 : tensor<2xui32>
    %1806 = mhlo.add %1801, %1805 : tensor<2xui32>
    %1807 = mhlo.shift_left %1805, %14 : tensor<2xui32>
    %1808 = mhlo.shift_right_logical %1805, %13 : tensor<2xui32>
    %1809 = mhlo.or %1807, %1808 : tensor<2xui32>
    %1810 = mhlo.xor %1806, %1809 : tensor<2xui32>
    %1811 = "mhlo.broadcast_in_dim"(%1761) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1812 = mhlo.add %1806, %1811 : tensor<2xui32>
    %1813 = "mhlo.broadcast_in_dim"(%1755) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1814 = mhlo.add %1810, %1813 : tensor<2xui32>
    %1815 = mhlo.add %1814, %18 : tensor<2xui32>
    %1816 = mhlo.add %1812, %1815 : tensor<2xui32>
    %1817 = mhlo.shift_left %1815, %11 : tensor<2xui32>
    %1818 = mhlo.shift_right_logical %1815, %10 : tensor<2xui32>
    %1819 = mhlo.or %1817, %1818 : tensor<2xui32>
    %1820 = mhlo.xor %1816, %1819 : tensor<2xui32>
    %1821 = mhlo.add %1816, %1820 : tensor<2xui32>
    %1822 = mhlo.shift_left %1820, %9 : tensor<2xui32>
    %1823 = mhlo.shift_right_logical %1820, %8 : tensor<2xui32>
    %1824 = mhlo.or %1822, %1823 : tensor<2xui32>
    %1825 = mhlo.xor %1821, %1824 : tensor<2xui32>
    %1826 = mhlo.add %1821, %1825 : tensor<2xui32>
    %1827 = mhlo.shift_left %1825, %6 : tensor<2xui32>
    %1828 = mhlo.shift_right_logical %1825, %7 : tensor<2xui32>
    %1829 = mhlo.or %1827, %1828 : tensor<2xui32>
    %1830 = mhlo.xor %1826, %1829 : tensor<2xui32>
    %1831 = mhlo.add %1826, %1830 : tensor<2xui32>
    %1832 = mhlo.shift_left %1830, %7 : tensor<2xui32>
    %1833 = mhlo.shift_right_logical %1830, %6 : tensor<2xui32>
    %1834 = mhlo.or %1832, %1833 : tensor<2xui32>
    %1835 = mhlo.xor %1831, %1834 : tensor<2xui32>
    %1836 = "mhlo.broadcast_in_dim"(%1755) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1837 = mhlo.add %1831, %1836 : tensor<2xui32>
    %1838 = "mhlo.broadcast_in_dim"(%1757) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1839 = mhlo.add %1835, %1838 : tensor<2xui32>
    %1840 = mhlo.add %1839, %16 : tensor<2xui32>
    %1841 = mhlo.add %1837, %1840 : tensor<2xui32>
    %1842 = mhlo.shift_left %1840, %8 : tensor<2xui32>
    %1843 = mhlo.shift_right_logical %1840, %9 : tensor<2xui32>
    %1844 = mhlo.or %1842, %1843 : tensor<2xui32>
    %1845 = mhlo.xor %1841, %1844 : tensor<2xui32>
    %1846 = mhlo.add %1841, %1845 : tensor<2xui32>
    %1847 = mhlo.shift_left %1845, %17 : tensor<2xui32>
    %1848 = mhlo.shift_right_logical %1845, %16 : tensor<2xui32>
    %1849 = mhlo.or %1847, %1848 : tensor<2xui32>
    %1850 = mhlo.xor %1846, %1849 : tensor<2xui32>
    %1851 = mhlo.add %1846, %1850 : tensor<2xui32>
    %1852 = mhlo.shift_left %1850, %15 : tensor<2xui32>
    %1853 = mhlo.shift_right_logical %1850, %15 : tensor<2xui32>
    %1854 = mhlo.or %1852, %1853 : tensor<2xui32>
    %1855 = mhlo.xor %1851, %1854 : tensor<2xui32>
    %1856 = mhlo.add %1851, %1855 : tensor<2xui32>
    %1857 = mhlo.shift_left %1855, %14 : tensor<2xui32>
    %1858 = mhlo.shift_right_logical %1855, %13 : tensor<2xui32>
    %1859 = mhlo.or %1857, %1858 : tensor<2xui32>
    %1860 = mhlo.xor %1856, %1859 : tensor<2xui32>
    %1861 = "mhlo.broadcast_in_dim"(%1757) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1862 = mhlo.add %1856, %1861 : tensor<2xui32>
    %1863 = "mhlo.broadcast_in_dim"(%1761) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1864 = mhlo.add %1860, %1863 : tensor<2xui32>
    %1865 = mhlo.add %1864, %12 : tensor<2xui32>
    %1866 = mhlo.add %1862, %1865 : tensor<2xui32>
    %1867 = mhlo.shift_left %1865, %11 : tensor<2xui32>
    %1868 = mhlo.shift_right_logical %1865, %10 : tensor<2xui32>
    %1869 = mhlo.or %1867, %1868 : tensor<2xui32>
    %1870 = mhlo.xor %1866, %1869 : tensor<2xui32>
    %1871 = mhlo.add %1866, %1870 : tensor<2xui32>
    %1872 = mhlo.shift_left %1870, %9 : tensor<2xui32>
    %1873 = mhlo.shift_right_logical %1870, %8 : tensor<2xui32>
    %1874 = mhlo.or %1872, %1873 : tensor<2xui32>
    %1875 = mhlo.xor %1871, %1874 : tensor<2xui32>
    %1876 = mhlo.add %1871, %1875 : tensor<2xui32>
    %1877 = mhlo.shift_left %1875, %6 : tensor<2xui32>
    %1878 = mhlo.shift_right_logical %1875, %7 : tensor<2xui32>
    %1879 = mhlo.or %1877, %1878 : tensor<2xui32>
    %1880 = mhlo.xor %1876, %1879 : tensor<2xui32>
    %1881 = mhlo.add %1876, %1880 : tensor<2xui32>
    %1882 = mhlo.shift_left %1880, %7 : tensor<2xui32>
    %1883 = mhlo.shift_right_logical %1880, %6 : tensor<2xui32>
    %1884 = mhlo.or %1882, %1883 : tensor<2xui32>
    %1885 = mhlo.xor %1881, %1884 : tensor<2xui32>
    %1886 = "mhlo.broadcast_in_dim"(%1761) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1887 = mhlo.add %1881, %1886 : tensor<2xui32>
    %1888 = "mhlo.broadcast_in_dim"(%1755) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<2xui32>
    %1889 = mhlo.add %1885, %1888 : tensor<2xui32>
    %1890 = mhlo.add %1889, %5 : tensor<2xui32>
    %1891 = "mhlo.concatenate"(%1887, %1890) {dimension = 0 : i64} : (tensor<2xui32>, tensor<2xui32>) -> tensor<4xui32>
    %1892 = "mhlo.reshape"(%1891) : (tensor<4xui32>) -> tensor<2x2xui32>
    %1893 = "mhlo.slice"(%1892) {limit_indices = dense<[1, 2]> : tensor<2xi64>, start_indices = dense<0> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<2x2xui32>) -> tensor<1x2xui32>
    %1894 = "mhlo.reshape"(%1893) : (tensor<1x2xui32>) -> tensor<2xui32>
    %1895 = "mhlo.slice"(%1892) {limit_indices = dense<2> : tensor<2xi64>, start_indices = dense<[1, 0]> : tensor<2xi64>, strides = dense<1> : tensor<2xi64>} : (tensor<2x2xui32>) -> tensor<1x2xui32>
    %1896 = "mhlo.reshape"(%1895) : (tensor<1x2xui32>) -> tensor<2xui32>
    %1897 = call @main_9(%68) : (tensor<f32>) -> tensor<f32>
    %1898 = call @main_10(%23) : (tensor<f32>) -> tensor<f32>
    %1899 = "mhlo.broadcast_in_dim"(%1897) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1x1xf32>
    %1900 = "mhlo.broadcast_in_dim"(%1898) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1x1xf32>
    %1901 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<10240xui32>
    %1902 = "mhlo.slice"(%1894) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xui32>) -> tensor<1xui32>
    %1903 = "mhlo.reshape"(%1902) : (tensor<1xui32>) -> tensor<ui32>
    %1904 = "mhlo.slice"(%1894) {limit_indices = dense<2> : tensor<1xi64>, start_indices = dense<1> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xui32>) -> tensor<1xui32>
    %1905 = "mhlo.reshape"(%1904) : (tensor<1xui32>) -> tensor<ui32>
    %1906 = "mhlo.slice"(%1901) {limit_indices = dense<5120> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<10240xui32>) -> tensor<5120xui32>
    %1907 = "mhlo.slice"(%1901) {limit_indices = dense<10240> : tensor<1xi64>, start_indices = dense<5120> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<10240xui32>) -> tensor<5120xui32>
    %1908 = mhlo.xor %1903, %1905 : tensor<ui32>
    %1909 = mhlo.xor %1908, %20 : tensor<ui32>
    %1910 = "mhlo.broadcast_in_dim"(%1903) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<5120xui32>
    %1911 = mhlo.add %1906, %1910 : tensor<5120xui32>
    %1912 = "mhlo.broadcast_in_dim"(%1905) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<5120xui32>
    %1913 = mhlo.add %1907, %1912 : tensor<5120xui32>
    %1914 = mhlo.add %1911, %1913 : tensor<5120xui32>
    %1915 = mhlo.shift_left %1913, %59 : tensor<5120xui32>
    %1916 = mhlo.shift_right_logical %1913, %58 : tensor<5120xui32>
    %1917 = mhlo.or %1915, %1916 : tensor<5120xui32>
    %1918 = mhlo.xor %1914, %1917 : tensor<5120xui32>
    %1919 = mhlo.add %1914, %1918 : tensor<5120xui32>
    %1920 = mhlo.shift_left %1918, %57 : tensor<5120xui32>
    %1921 = mhlo.shift_right_logical %1918, %56 : tensor<5120xui32>
    %1922 = mhlo.or %1920, %1921 : tensor<5120xui32>
    %1923 = mhlo.xor %1919, %1922 : tensor<5120xui32>
    %1924 = mhlo.add %1919, %1923 : tensor<5120xui32>
    %1925 = mhlo.shift_left %1923, %54 : tensor<5120xui32>
    %1926 = mhlo.shift_right_logical %1923, %55 : tensor<5120xui32>
    %1927 = mhlo.or %1925, %1926 : tensor<5120xui32>
    %1928 = mhlo.xor %1924, %1927 : tensor<5120xui32>
    %1929 = mhlo.add %1924, %1928 : tensor<5120xui32>
    %1930 = mhlo.shift_left %1928, %55 : tensor<5120xui32>
    %1931 = mhlo.shift_right_logical %1928, %54 : tensor<5120xui32>
    %1932 = mhlo.or %1930, %1931 : tensor<5120xui32>
    %1933 = mhlo.xor %1929, %1932 : tensor<5120xui32>
    %1934 = "mhlo.broadcast_in_dim"(%1905) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<5120xui32>
    %1935 = mhlo.add %1929, %1934 : tensor<5120xui32>
    %1936 = "mhlo.broadcast_in_dim"(%1909) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<5120xui32>
    %1937 = mhlo.add %1933, %1936 : tensor<5120xui32>
    %1938 = mhlo.add %1937, %67 : tensor<5120xui32>
    %1939 = mhlo.add %1935, %1938 : tensor<5120xui32>
    %1940 = mhlo.shift_left %1938, %56 : tensor<5120xui32>
    %1941 = mhlo.shift_right_logical %1938, %57 : tensor<5120xui32>
    %1942 = mhlo.or %1940, %1941 : tensor<5120xui32>
    %1943 = mhlo.xor %1939, %1942 : tensor<5120xui32>
    %1944 = mhlo.add %1939, %1943 : tensor<5120xui32>
    %1945 = mhlo.shift_left %1943, %65 : tensor<5120xui32>
    %1946 = mhlo.shift_right_logical %1943, %64 : tensor<5120xui32>
    %1947 = mhlo.or %1945, %1946 : tensor<5120xui32>
    %1948 = mhlo.xor %1944, %1947 : tensor<5120xui32>
    %1949 = mhlo.add %1944, %1948 : tensor<5120xui32>
    %1950 = mhlo.shift_left %1948, %63 : tensor<5120xui32>
    %1951 = mhlo.shift_right_logical %1948, %63 : tensor<5120xui32>
    %1952 = mhlo.or %1950, %1951 : tensor<5120xui32>
    %1953 = mhlo.xor %1949, %1952 : tensor<5120xui32>
    %1954 = mhlo.add %1949, %1953 : tensor<5120xui32>
    %1955 = mhlo.shift_left %1953, %62 : tensor<5120xui32>
    %1956 = mhlo.shift_right_logical %1953, %61 : tensor<5120xui32>
    %1957 = mhlo.or %1955, %1956 : tensor<5120xui32>
    %1958 = mhlo.xor %1954, %1957 : tensor<5120xui32>
    %1959 = "mhlo.broadcast_in_dim"(%1909) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<5120xui32>
    %1960 = mhlo.add %1954, %1959 : tensor<5120xui32>
    %1961 = "mhlo.broadcast_in_dim"(%1903) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<5120xui32>
    %1962 = mhlo.add %1958, %1961 : tensor<5120xui32>
    %1963 = mhlo.add %1962, %66 : tensor<5120xui32>
    %1964 = mhlo.add %1960, %1963 : tensor<5120xui32>
    %1965 = mhlo.shift_left %1963, %59 : tensor<5120xui32>
    %1966 = mhlo.shift_right_logical %1963, %58 : tensor<5120xui32>
    %1967 = mhlo.or %1965, %1966 : tensor<5120xui32>
    %1968 = mhlo.xor %1964, %1967 : tensor<5120xui32>
    %1969 = mhlo.add %1964, %1968 : tensor<5120xui32>
    %1970 = mhlo.shift_left %1968, %57 : tensor<5120xui32>
    %1971 = mhlo.shift_right_logical %1968, %56 : tensor<5120xui32>
    %1972 = mhlo.or %1970, %1971 : tensor<5120xui32>
    %1973 = mhlo.xor %1969, %1972 : tensor<5120xui32>
    %1974 = mhlo.add %1969, %1973 : tensor<5120xui32>
    %1975 = mhlo.shift_left %1973, %54 : tensor<5120xui32>
    %1976 = mhlo.shift_right_logical %1973, %55 : tensor<5120xui32>
    %1977 = mhlo.or %1975, %1976 : tensor<5120xui32>
    %1978 = mhlo.xor %1974, %1977 : tensor<5120xui32>
    %1979 = mhlo.add %1974, %1978 : tensor<5120xui32>
    %1980 = mhlo.shift_left %1978, %55 : tensor<5120xui32>
    %1981 = mhlo.shift_right_logical %1978, %54 : tensor<5120xui32>
    %1982 = mhlo.or %1980, %1981 : tensor<5120xui32>
    %1983 = mhlo.xor %1979, %1982 : tensor<5120xui32>
    %1984 = "mhlo.broadcast_in_dim"(%1903) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<5120xui32>
    %1985 = mhlo.add %1979, %1984 : tensor<5120xui32>
    %1986 = "mhlo.broadcast_in_dim"(%1905) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<5120xui32>
    %1987 = mhlo.add %1983, %1986 : tensor<5120xui32>
    %1988 = mhlo.add %1987, %64 : tensor<5120xui32>
    %1989 = mhlo.add %1985, %1988 : tensor<5120xui32>
    %1990 = mhlo.shift_left %1988, %56 : tensor<5120xui32>
    %1991 = mhlo.shift_right_logical %1988, %57 : tensor<5120xui32>
    %1992 = mhlo.or %1990, %1991 : tensor<5120xui32>
    %1993 = mhlo.xor %1989, %1992 : tensor<5120xui32>
    %1994 = mhlo.add %1989, %1993 : tensor<5120xui32>
    %1995 = mhlo.shift_left %1993, %65 : tensor<5120xui32>
    %1996 = mhlo.shift_right_logical %1993, %64 : tensor<5120xui32>
    %1997 = mhlo.or %1995, %1996 : tensor<5120xui32>
    %1998 = mhlo.xor %1994, %1997 : tensor<5120xui32>
    %1999 = mhlo.add %1994, %1998 : tensor<5120xui32>
    %2000 = mhlo.shift_left %1998, %63 : tensor<5120xui32>
    %2001 = mhlo.shift_right_logical %1998, %63 : tensor<5120xui32>
    %2002 = mhlo.or %2000, %2001 : tensor<5120xui32>
    %2003 = mhlo.xor %1999, %2002 : tensor<5120xui32>
    %2004 = mhlo.add %1999, %2003 : tensor<5120xui32>
    %2005 = mhlo.shift_left %2003, %62 : tensor<5120xui32>
    %2006 = mhlo.shift_right_logical %2003, %61 : tensor<5120xui32>
    %2007 = mhlo.or %2005, %2006 : tensor<5120xui32>
    %2008 = mhlo.xor %2004, %2007 : tensor<5120xui32>
    %2009 = "mhlo.broadcast_in_dim"(%1905) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<5120xui32>
    %2010 = mhlo.add %2004, %2009 : tensor<5120xui32>
    %2011 = "mhlo.broadcast_in_dim"(%1909) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<5120xui32>
    %2012 = mhlo.add %2008, %2011 : tensor<5120xui32>
    %2013 = mhlo.add %2012, %60 : tensor<5120xui32>
    %2014 = mhlo.add %2010, %2013 : tensor<5120xui32>
    %2015 = mhlo.shift_left %2013, %59 : tensor<5120xui32>
    %2016 = mhlo.shift_right_logical %2013, %58 : tensor<5120xui32>
    %2017 = mhlo.or %2015, %2016 : tensor<5120xui32>
    %2018 = mhlo.xor %2014, %2017 : tensor<5120xui32>
    %2019 = mhlo.add %2014, %2018 : tensor<5120xui32>
    %2020 = mhlo.shift_left %2018, %57 : tensor<5120xui32>
    %2021 = mhlo.shift_right_logical %2018, %56 : tensor<5120xui32>
    %2022 = mhlo.or %2020, %2021 : tensor<5120xui32>
    %2023 = mhlo.xor %2019, %2022 : tensor<5120xui32>
    %2024 = mhlo.add %2019, %2023 : tensor<5120xui32>
    %2025 = mhlo.shift_left %2023, %54 : tensor<5120xui32>
    %2026 = mhlo.shift_right_logical %2023, %55 : tensor<5120xui32>
    %2027 = mhlo.or %2025, %2026 : tensor<5120xui32>
    %2028 = mhlo.xor %2024, %2027 : tensor<5120xui32>
    %2029 = mhlo.add %2024, %2028 : tensor<5120xui32>
    %2030 = mhlo.shift_left %2028, %55 : tensor<5120xui32>
    %2031 = mhlo.shift_right_logical %2028, %54 : tensor<5120xui32>
    %2032 = mhlo.or %2030, %2031 : tensor<5120xui32>
    %2033 = mhlo.xor %2029, %2032 : tensor<5120xui32>
    %2034 = "mhlo.broadcast_in_dim"(%1909) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<5120xui32>
    %2035 = mhlo.add %2029, %2034 : tensor<5120xui32>
    %2036 = "mhlo.broadcast_in_dim"(%1903) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<5120xui32>
    %2037 = mhlo.add %2033, %2036 : tensor<5120xui32>
    %2038 = mhlo.add %2037, %53 : tensor<5120xui32>
    %2039 = "mhlo.concatenate"(%2035, %2038) {dimension = 0 : i64} : (tensor<5120xui32>, tensor<5120xui32>) -> tensor<10240xui32>
    %2040 = "mhlo.reshape"(%2039) : (tensor<10240xui32>) -> tensor<1024x10xui32>
    %2041 = mhlo.shift_right_logical %2040, %52 : tensor<1024x10xui32>
    %2042 = mhlo.or %2041, %51 : tensor<1024x10xui32>
    %2043 = "mhlo.bitcast_convert"(%2042) : (tensor<1024x10xui32>) -> tensor<1024x10xf32>
    %2044 = mhlo.subtract %2043, %50 : tensor<1024x10xf32>
    %2045 = mhlo.subtract %1900, %1899 : tensor<1x1xf32>
    %2046 = "mhlo.broadcast_in_dim"(%2045) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor<1x1xf32>) -> tensor<1024x10xf32>
    %2047 = mhlo.multiply %2044, %2046 : tensor<1024x10xf32>
    %2048 = "mhlo.broadcast_in_dim"(%1899) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor<1x1xf32>) -> tensor<1024x10xf32>
    %2049 = mhlo.add %2047, %2048 : tensor<1024x10xf32>
    %2050 = "mhlo.broadcast_in_dim"(%1899) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor<1x1xf32>) -> tensor<1024x10xf32>
    %2051 = mhlo.maximum %2050, %2049 : tensor<1024x10xf32>
    %2052 = call @main_11(%2051) : (tensor<1024x10xf32>) -> tensor<1024x10xf32>
    %2053 = mhlo.multiply %2052, %49 : tensor<1024x10xf32>
    %2054 = chlo.next_after %48, %47 : tensor<f32>, tensor<f32> -> tensor<f32>
    %2055 = chlo.next_after %46, %45 : tensor<f32>, tensor<f32> -> tensor<f32>
    %2056 = call @jit_clip_12(%2053, %2054, %2055) : (tensor<1024x10xf32>, tensor<f32>, tensor<f32>) -> tensor<1024x10xf32>
    %2057 = mhlo.multiply %2056, %44 : tensor<1024x10xf32>
    %2058 = "mhlo.iota"() {iota_dimension = 0 : i64} : () -> tensor<10xui32>
    %2059 = "mhlo.slice"(%1896) {limit_indices = dense<1> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xui32>) -> tensor<1xui32>
    %2060 = "mhlo.reshape"(%2059) : (tensor<1xui32>) -> tensor<ui32>
    %2061 = "mhlo.slice"(%1896) {limit_indices = dense<2> : tensor<1xi64>, start_indices = dense<1> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<2xui32>) -> tensor<1xui32>
    %2062 = "mhlo.reshape"(%2061) : (tensor<1xui32>) -> tensor<ui32>
    %2063 = "mhlo.slice"(%2058) {limit_indices = dense<5> : tensor<1xi64>, start_indices = dense<0> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<10xui32>) -> tensor<5xui32>
    %2064 = "mhlo.slice"(%2058) {limit_indices = dense<10> : tensor<1xi64>, start_indices = dense<5> : tensor<1xi64>, strides = dense<1> : tensor<1xi64>} : (tensor<10xui32>) -> tensor<5xui32>
    %2065 = mhlo.xor %2060, %2062 : tensor<ui32>
    %2066 = mhlo.xor %2065, %20 : tensor<ui32>
    %2067 = "mhlo.broadcast_in_dim"(%2060) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<5xui32>
    %2068 = mhlo.add %2063, %2067 : tensor<5xui32>
    %2069 = "mhlo.broadcast_in_dim"(%2062) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<5xui32>
    %2070 = mhlo.add %2064, %2069 : tensor<5xui32>
    %2071 = mhlo.add %2068, %2070 : tensor<5xui32>
    %2072 = mhlo.shift_left %2070, %35 : tensor<5xui32>
    %2073 = mhlo.shift_right_logical %2070, %34 : tensor<5xui32>
    %2074 = mhlo.or %2072, %2073 : tensor<5xui32>
    %2075 = mhlo.xor %2071, %2074 : tensor<5xui32>
    %2076 = mhlo.add %2071, %2075 : tensor<5xui32>
    %2077 = mhlo.shift_left %2075, %33 : tensor<5xui32>
    %2078 = mhlo.shift_right_logical %2075, %32 : tensor<5xui32>
    %2079 = mhlo.or %2077, %2078 : tensor<5xui32>
    %2080 = mhlo.xor %2076, %2079 : tensor<5xui32>
    %2081 = mhlo.add %2076, %2080 : tensor<5xui32>
    %2082 = mhlo.shift_left %2080, %30 : tensor<5xui32>
    %2083 = mhlo.shift_right_logical %2080, %31 : tensor<5xui32>
    %2084 = mhlo.or %2082, %2083 : tensor<5xui32>
    %2085 = mhlo.xor %2081, %2084 : tensor<5xui32>
    %2086 = mhlo.add %2081, %2085 : tensor<5xui32>
    %2087 = mhlo.shift_left %2085, %31 : tensor<5xui32>
    %2088 = mhlo.shift_right_logical %2085, %30 : tensor<5xui32>
    %2089 = mhlo.or %2087, %2088 : tensor<5xui32>
    %2090 = mhlo.xor %2086, %2089 : tensor<5xui32>
    %2091 = "mhlo.broadcast_in_dim"(%2062) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<5xui32>
    %2092 = mhlo.add %2086, %2091 : tensor<5xui32>
    %2093 = "mhlo.broadcast_in_dim"(%2066) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<5xui32>
    %2094 = mhlo.add %2090, %2093 : tensor<5xui32>
    %2095 = mhlo.add %2094, %43 : tensor<5xui32>
    %2096 = mhlo.add %2092, %2095 : tensor<5xui32>
    %2097 = mhlo.shift_left %2095, %32 : tensor<5xui32>
    %2098 = mhlo.shift_right_logical %2095, %33 : tensor<5xui32>
    %2099 = mhlo.or %2097, %2098 : tensor<5xui32>
    %2100 = mhlo.xor %2096, %2099 : tensor<5xui32>
    %2101 = mhlo.add %2096, %2100 : tensor<5xui32>
    %2102 = mhlo.shift_left %2100, %41 : tensor<5xui32>
    %2103 = mhlo.shift_right_logical %2100, %40 : tensor<5xui32>
    %2104 = mhlo.or %2102, %2103 : tensor<5xui32>
    %2105 = mhlo.xor %2101, %2104 : tensor<5xui32>
    %2106 = mhlo.add %2101, %2105 : tensor<5xui32>
    %2107 = mhlo.shift_left %2105, %39 : tensor<5xui32>
    %2108 = mhlo.shift_right_logical %2105, %39 : tensor<5xui32>
    %2109 = mhlo.or %2107, %2108 : tensor<5xui32>
    %2110 = mhlo.xor %2106, %2109 : tensor<5xui32>
    %2111 = mhlo.add %2106, %2110 : tensor<5xui32>
    %2112 = mhlo.shift_left %2110, %38 : tensor<5xui32>
    %2113 = mhlo.shift_right_logical %2110, %37 : tensor<5xui32>
    %2114 = mhlo.or %2112, %2113 : tensor<5xui32>
    %2115 = mhlo.xor %2111, %2114 : tensor<5xui32>
    %2116 = "mhlo.broadcast_in_dim"(%2066) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<5xui32>
    %2117 = mhlo.add %2111, %2116 : tensor<5xui32>
    %2118 = "mhlo.broadcast_in_dim"(%2060) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<5xui32>
    %2119 = mhlo.add %2115, %2118 : tensor<5xui32>
    %2120 = mhlo.add %2119, %42 : tensor<5xui32>
    %2121 = mhlo.add %2117, %2120 : tensor<5xui32>
    %2122 = mhlo.shift_left %2120, %35 : tensor<5xui32>
    %2123 = mhlo.shift_right_logical %2120, %34 : tensor<5xui32>
    %2124 = mhlo.or %2122, %2123 : tensor<5xui32>
    %2125 = mhlo.xor %2121, %2124 : tensor<5xui32>
    %2126 = mhlo.add %2121, %2125 : tensor<5xui32>
    %2127 = mhlo.shift_left %2125, %33 : tensor<5xui32>
    %2128 = mhlo.shift_right_logical %2125, %32 : tensor<5xui32>
    %2129 = mhlo.or %2127, %2128 : tensor<5xui32>
    %2130 = mhlo.xor %2126, %2129 : tensor<5xui32>
    %2131 = mhlo.add %2126, %2130 : tensor<5xui32>
    %2132 = mhlo.shift_left %2130, %30 : tensor<5xui32>
    %2133 = mhlo.shift_right_logical %2130, %31 : tensor<5xui32>
    %2134 = mhlo.or %2132, %2133 : tensor<5xui32>
    %2135 = mhlo.xor %2131, %2134 : tensor<5xui32>
    %2136 = mhlo.add %2131, %2135 : tensor<5xui32>
    %2137 = mhlo.shift_left %2135, %31 : tensor<5xui32>
    %2138 = mhlo.shift_right_logical %2135, %30 : tensor<5xui32>
    %2139 = mhlo.or %2137, %2138 : tensor<5xui32>
    %2140 = mhlo.xor %2136, %2139 : tensor<5xui32>
    %2141 = "mhlo.broadcast_in_dim"(%2060) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<5xui32>
    %2142 = mhlo.add %2136, %2141 : tensor<5xui32>
    %2143 = "mhlo.broadcast_in_dim"(%2062) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<5xui32>
    %2144 = mhlo.add %2140, %2143 : tensor<5xui32>
    %2145 = mhlo.add %2144, %40 : tensor<5xui32>
    %2146 = mhlo.add %2142, %2145 : tensor<5xui32>
    %2147 = mhlo.shift_left %2145, %32 : tensor<5xui32>
    %2148 = mhlo.shift_right_logical %2145, %33 : tensor<5xui32>
    %2149 = mhlo.or %2147, %2148 : tensor<5xui32>
    %2150 = mhlo.xor %2146, %2149 : tensor<5xui32>
    %2151 = mhlo.add %2146, %2150 : tensor<5xui32>
    %2152 = mhlo.shift_left %2150, %41 : tensor<5xui32>
    %2153 = mhlo.shift_right_logical %2150, %40 : tensor<5xui32>
    %2154 = mhlo.or %2152, %2153 : tensor<5xui32>
    %2155 = mhlo.xor %2151, %2154 : tensor<5xui32>
    %2156 = mhlo.add %2151, %2155 : tensor<5xui32>
    %2157 = mhlo.shift_left %2155, %39 : tensor<5xui32>
    %2158 = mhlo.shift_right_logical %2155, %39 : tensor<5xui32>
    %2159 = mhlo.or %2157, %2158 : tensor<5xui32>
    %2160 = mhlo.xor %2156, %2159 : tensor<5xui32>
    %2161 = mhlo.add %2156, %2160 : tensor<5xui32>
    %2162 = mhlo.shift_left %2160, %38 : tensor<5xui32>
    %2163 = mhlo.shift_right_logical %2160, %37 : tensor<5xui32>
    %2164 = mhlo.or %2162, %2163 : tensor<5xui32>
    %2165 = mhlo.xor %2161, %2164 : tensor<5xui32>
    %2166 = "mhlo.broadcast_in_dim"(%2062) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<5xui32>
    %2167 = mhlo.add %2161, %2166 : tensor<5xui32>
    %2168 = "mhlo.broadcast_in_dim"(%2066) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<5xui32>
    %2169 = mhlo.add %2165, %2168 : tensor<5xui32>
    %2170 = mhlo.add %2169, %36 : tensor<5xui32>
    %2171 = mhlo.add %2167, %2170 : tensor<5xui32>
    %2172 = mhlo.shift_left %2170, %35 : tensor<5xui32>
    %2173 = mhlo.shift_right_logical %2170, %34 : tensor<5xui32>
    %2174 = mhlo.or %2172, %2173 : tensor<5xui32>
    %2175 = mhlo.xor %2171, %2174 : tensor<5xui32>
    %2176 = mhlo.add %2171, %2175 : tensor<5xui32>
    %2177 = mhlo.shift_left %2175, %33 : tensor<5xui32>
    %2178 = mhlo.shift_right_logical %2175, %32 : tensor<5xui32>
    %2179 = mhlo.or %2177, %2178 : tensor<5xui32>
    %2180 = mhlo.xor %2176, %2179 : tensor<5xui32>
    %2181 = mhlo.add %2176, %2180 : tensor<5xui32>
    %2182 = mhlo.shift_left %2180, %30 : tensor<5xui32>
    %2183 = mhlo.shift_right_logical %2180, %31 : tensor<5xui32>
    %2184 = mhlo.or %2182, %2183 : tensor<5xui32>
    %2185 = mhlo.xor %2181, %2184 : tensor<5xui32>
    %2186 = mhlo.add %2181, %2185 : tensor<5xui32>
    %2187 = mhlo.shift_left %2185, %31 : tensor<5xui32>
    %2188 = mhlo.shift_right_logical %2185, %30 : tensor<5xui32>
    %2189 = mhlo.or %2187, %2188 : tensor<5xui32>
    %2190 = mhlo.xor %2186, %2189 : tensor<5xui32>
    %2191 = "mhlo.broadcast_in_dim"(%2066) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<5xui32>
    %2192 = mhlo.add %2186, %2191 : tensor<5xui32>
    %2193 = "mhlo.broadcast_in_dim"(%2060) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<ui32>) -> tensor<5xui32>
    %2194 = mhlo.add %2190, %2193 : tensor<5xui32>
    %2195 = mhlo.add %2194, %29 : tensor<5xui32>
    %2196 = "mhlo.concatenate"(%2192, %2195) {dimension = 0 : i64} : (tensor<5xui32>, tensor<5xui32>) -> tensor<10xui32>
    %2197 = mhlo.shift_right_logical %2196, %28 : tensor<10xui32>
    %2198 = mhlo.or %2197, %27 : tensor<10xui32>
    %2199 = "mhlo.bitcast_convert"(%2198) : (tensor<10xui32>) -> tensor<10xf32>
    %2200 = mhlo.subtract %2199, %26 : tensor<10xf32>
    %2201 = mhlo.multiply %2200, %25 : tensor<10xf32>
    %2202 = mhlo.add %2201, %24 : tensor<10xf32>
    %2203 = mhlo.maximum %2202, %24 : tensor<10xf32>
    %2204 = call @main_13(%2203) : (tensor<10xf32>) -> tensor<10xf32>
    %2205 = mhlo.multiply %2204, %22 : tensor<10xf32>
    %2206 = mhlo.multiply %2205, %21 : tensor<10xf32>
    return %579, %4, %728, %2, %1319, %3, %1468, %2, %2057, %1, %2206, %0 : tensor<784x1024xf32>, tensor<784x1024xf32>, tensor<1024xf32>, tensor<1024xf32>, tensor<1024x1024xf32>, tensor<1024x1024xf32>, tensor<1024xf32>, tensor<1024xf32>, tensor<1024x10xf32>, tensor<1024x10xf32>, tensor<10xf32>, tensor<10xf32>
  }
  func private @main_0(%arg0: tensor<f32>) -> tensor<f32> {
    %0 = mhlo.constant dense<-0.0142647391> : tensor<f32>
    %1 = mhlo.constant dense<-0.00737332925> : tensor<f32>
    %2 = mhlo.constant dense<-0.00168282702> : tensor<f32>
    %3 = mhlo.constant dense<-2.13374049E-4> : tensor<f32>
    %4 = mhlo.constant dense<-1.45660715E-5> : tensor<f32>
    %5 = mhlo.constant dense<0.000000e+00> : tensor<f32>
    %6 = mhlo.constant dense<-0.0160960332> : tensor<f32>
    %7 = mhlo.constant dense<-2.954600e-03> : tensor<f32>
    %8 = mhlo.constant dense<-7.34990637E-4> : tensor<f32>
    %9 = mhlo.constant dense<-5.69250624E-5> : tensor<f32>
    %10 = mhlo.constant dense<-2.10102394E-6> : tensor<f32>
    %11 = mhlo.constant dense<2.77068146E-8> : tensor<f32>
    %12 = mhlo.constant dense<-2.72614237E-10> : tensor<f32>
    %13 = mhlo.constant dense<4.000000e+00> : tensor<f32>
    %14 = mhlo.constant dense<-4.000000e+00> : tensor<f32>
    %15 = "mhlo.clamp"(%14, %arg0, %13) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
    %16 = mhlo.multiply %15, %15 : tensor<f32>
    %17 = mhlo.multiply %16, %5 : tensor<f32>
    %18 = mhlo.add %17, %12 : tensor<f32>
    %19 = mhlo.multiply %18, %16 : tensor<f32>
    %20 = mhlo.add %19, %11 : tensor<f32>
    %21 = mhlo.multiply %20, %16 : tensor<f32>
    %22 = mhlo.add %21, %10 : tensor<f32>
    %23 = mhlo.multiply %22, %16 : tensor<f32>
    %24 = mhlo.add %23, %9 : tensor<f32>
    %25 = mhlo.multiply %24, %16 : tensor<f32>
    %26 = mhlo.add %25, %8 : tensor<f32>
    %27 = mhlo.multiply %26, %16 : tensor<f32>
    %28 = mhlo.add %27, %7 : tensor<f32>
    %29 = mhlo.multiply %28, %16 : tensor<f32>
    %30 = mhlo.add %29, %6 : tensor<f32>
    %31 = mhlo.multiply %15, %30 : tensor<f32>
    %32 = mhlo.multiply %16, %5 : tensor<f32>
    %33 = mhlo.add %32, %4 : tensor<f32>
    %34 = mhlo.multiply %33, %16 : tensor<f32>
    %35 = mhlo.add %34, %3 : tensor<f32>
    %36 = mhlo.multiply %35, %16 : tensor<f32>
    %37 = mhlo.add %36, %2 : tensor<f32>
    %38 = mhlo.multiply %37, %16 : tensor<f32>
    %39 = mhlo.add %38, %1 : tensor<f32>
    %40 = mhlo.multiply %39, %16 : tensor<f32>
    %41 = mhlo.add %40, %0 : tensor<f32>
    %42 = mhlo.divide %31, %41 : tensor<f32>
    return %42 : tensor<f32>
  }
  func private @main_1(%arg0: tensor<f32>) -> tensor<f32> {
    %0 = mhlo.constant dense<-0.0142647391> : tensor<f32>
    %1 = mhlo.constant dense<-0.00737332925> : tensor<f32>
    %2 = mhlo.constant dense<-0.00168282702> : tensor<f32>
    %3 = mhlo.constant dense<-2.13374049E-4> : tensor<f32>
    %4 = mhlo.constant dense<-1.45660715E-5> : tensor<f32>
    %5 = mhlo.constant dense<0.000000e+00> : tensor<f32>
    %6 = mhlo.constant dense<-0.0160960332> : tensor<f32>
    %7 = mhlo.constant dense<-2.954600e-03> : tensor<f32>
    %8 = mhlo.constant dense<-7.34990637E-4> : tensor<f32>
    %9 = mhlo.constant dense<-5.69250624E-5> : tensor<f32>
    %10 = mhlo.constant dense<-2.10102394E-6> : tensor<f32>
    %11 = mhlo.constant dense<2.77068146E-8> : tensor<f32>
    %12 = mhlo.constant dense<-2.72614237E-10> : tensor<f32>
    %13 = mhlo.constant dense<4.000000e+00> : tensor<f32>
    %14 = mhlo.constant dense<-4.000000e+00> : tensor<f32>
    %15 = "mhlo.clamp"(%14, %arg0, %13) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
    %16 = mhlo.multiply %15, %15 : tensor<f32>
    %17 = mhlo.multiply %16, %5 : tensor<f32>
    %18 = mhlo.add %17, %12 : tensor<f32>
    %19 = mhlo.multiply %18, %16 : tensor<f32>
    %20 = mhlo.add %19, %11 : tensor<f32>
    %21 = mhlo.multiply %20, %16 : tensor<f32>
    %22 = mhlo.add %21, %10 : tensor<f32>
    %23 = mhlo.multiply %22, %16 : tensor<f32>
    %24 = mhlo.add %23, %9 : tensor<f32>
    %25 = mhlo.multiply %24, %16 : tensor<f32>
    %26 = mhlo.add %25, %8 : tensor<f32>
    %27 = mhlo.multiply %26, %16 : tensor<f32>
    %28 = mhlo.add %27, %7 : tensor<f32>
    %29 = mhlo.multiply %28, %16 : tensor<f32>
    %30 = mhlo.add %29, %6 : tensor<f32>
    %31 = mhlo.multiply %15, %30 : tensor<f32>
    %32 = mhlo.multiply %16, %5 : tensor<f32>
    %33 = mhlo.add %32, %4 : tensor<f32>
    %34 = mhlo.multiply %33, %16 : tensor<f32>
    %35 = mhlo.add %34, %3 : tensor<f32>
    %36 = mhlo.multiply %35, %16 : tensor<f32>
    %37 = mhlo.add %36, %2 : tensor<f32>
    %38 = mhlo.multiply %37, %16 : tensor<f32>
    %39 = mhlo.add %38, %1 : tensor<f32>
    %40 = mhlo.multiply %39, %16 : tensor<f32>
    %41 = mhlo.add %40, %0 : tensor<f32>
    %42 = mhlo.divide %31, %41 : tensor<f32>
    return %42 : tensor<f32>
  }
  func private @main_2(%arg0: tensor<784x1024xf32>) -> tensor<784x1024xf32> {
    %0 = mhlo.constant dense<3.000000e+00> : tensor<784x1024xf32>
    %1 = mhlo.constant dense<2.500000e+00> : tensor<784x1024xf32>
    %2 = mhlo.constant dense<-2.00214257E-4> : tensor<784x1024xf32>
    %3 = mhlo.constant dense<2.81022636E-8> : tensor<784x1024xf32>
    %4 = mhlo.constant dense<1.00950558E-4> : tensor<784x1024xf32>
    %5 = mhlo.constant dense<3.43273939E-7> : tensor<784x1024xf32>
    %6 = mhlo.constant dense<0.00134934322> : tensor<784x1024xf32>
    %7 = mhlo.constant dense<-3.5233877E-6> : tensor<784x1024xf32>
    %8 = mhlo.constant dense<-0.00367342844> : tensor<784x1024xf32>
    %9 = mhlo.constant dense<-4.39150654E-6> : tensor<784x1024xf32>
    %10 = mhlo.constant dense<0.00573950773> : tensor<784x1024xf32>
    %11 = mhlo.constant dense<2.1858087E-4> : tensor<784x1024xf32>
    %12 = mhlo.constant dense<-0.0076224613> : tensor<784x1024xf32>
    %13 = mhlo.constant dense<-0.00125372503> : tensor<784x1024xf32>
    %14 = mhlo.constant dense<0.00943887047> : tensor<784x1024xf32>
    %15 = mhlo.constant dense<-0.00417768164> : tensor<784x1024xf32>
    %16 = mhlo.constant dense<1.00167406> : tensor<784x1024xf32>
    %17 = mhlo.constant dense<0.246640727> : tensor<784x1024xf32>
    %18 = mhlo.constant dense<2.83297682> : tensor<784x1024xf32>
    %19 = mhlo.constant dense<1.50140941> : tensor<784x1024xf32>
    %20 = mhlo.constant dense<5.000000e+00> : tensor<784x1024xf32>
    %21 = mhlo.constant dense<0x7F800000> : tensor<784x1024xf32>
    %22 = mhlo.constant dense<1.000000e+00> : tensor<784x1024xf32>
    %23 = "mhlo.abs"(%arg0) : (tensor<784x1024xf32>) -> tensor<784x1024xf32>
    %24 = "mhlo.compare"(%23, %22) {comparison_direction = "EQ"} : (tensor<784x1024xf32>, tensor<784x1024xf32>) -> tensor<784x1024xi1>
    %25 = mhlo.multiply %arg0, %21 : tensor<784x1024xf32>
    %26 = "mhlo.negate"(%arg0) : (tensor<784x1024xf32>) -> tensor<784x1024xf32>
    %27 = mhlo.multiply %26, %arg0 : tensor<784x1024xf32>
    %28 = "mhlo.log_plus_one"(%27) : (tensor<784x1024xf32>) -> tensor<784x1024xf32>
    %29 = "mhlo.negate"(%28) : (tensor<784x1024xf32>) -> tensor<784x1024xf32>
    %30 = "mhlo.compare"(%29, %20) {comparison_direction = "LT"} : (tensor<784x1024xf32>, tensor<784x1024xf32>) -> tensor<784x1024xi1>
    %31 = "mhlo.select"(%30, %19, %18) : (tensor<784x1024xi1>, tensor<784x1024xf32>, tensor<784x1024xf32>) -> tensor<784x1024xf32>
    %32 = "mhlo.select"(%30, %17, %16) : (tensor<784x1024xi1>, tensor<784x1024xf32>, tensor<784x1024xf32>) -> tensor<784x1024xf32>
    %33 = "mhlo.select"(%30, %15, %14) : (tensor<784x1024xi1>, tensor<784x1024xf32>, tensor<784x1024xf32>) -> tensor<784x1024xf32>
    %34 = "mhlo.select"(%30, %13, %12) : (tensor<784x1024xi1>, tensor<784x1024xf32>, tensor<784x1024xf32>) -> tensor<784x1024xf32>
    %35 = "mhlo.select"(%30, %11, %10) : (tensor<784x1024xi1>, tensor<784x1024xf32>, tensor<784x1024xf32>) -> tensor<784x1024xf32>
    %36 = "mhlo.select"(%30, %9, %8) : (tensor<784x1024xi1>, tensor<784x1024xf32>, tensor<784x1024xf32>) -> tensor<784x1024xf32>
    %37 = "mhlo.select"(%30, %7, %6) : (tensor<784x1024xi1>, tensor<784x1024xf32>, tensor<784x1024xf32>) -> tensor<784x1024xf32>
    %38 = "mhlo.select"(%30, %5, %4) : (tensor<784x1024xi1>, tensor<784x1024xf32>, tensor<784x1024xf32>) -> tensor<784x1024xf32>
    %39 = "mhlo.select"(%30, %3, %2) : (tensor<784x1024xi1>, tensor<784x1024xf32>, tensor<784x1024xf32>) -> tensor<784x1024xf32>
    %40 = mhlo.subtract %29, %1 : tensor<784x1024xf32>
    %41 = "mhlo.sqrt"(%29) : (tensor<784x1024xf32>) -> tensor<784x1024xf32>
    %42 = mhlo.subtract %41, %0 : tensor<784x1024xf32>
    %43 = "mhlo.select"(%30, %40, %42) : (tensor<784x1024xi1>, tensor<784x1024xf32>, tensor<784x1024xf32>) -> tensor<784x1024xf32>
    %44 = mhlo.multiply %39, %43 : tensor<784x1024xf32>
    %45 = mhlo.add %38, %44 : tensor<784x1024xf32>
    %46 = mhlo.multiply %45, %43 : tensor<784x1024xf32>
    %47 = mhlo.add %37, %46 : tensor<784x1024xf32>
    %48 = mhlo.multiply %47, %43 : tensor<784x1024xf32>
    %49 = mhlo.add %36, %48 : tensor<784x1024xf32>
    %50 = mhlo.multiply %49, %43 : tensor<784x1024xf32>
    %51 = mhlo.add %35, %50 : tensor<784x1024xf32>
    %52 = mhlo.multiply %51, %43 : tensor<784x1024xf32>
    %53 = mhlo.add %34, %52 : tensor<784x1024xf32>
    %54 = mhlo.multiply %53, %43 : tensor<784x1024xf32>
    %55 = mhlo.add %33, %54 : tensor<784x1024xf32>
    %56 = mhlo.multiply %55, %43 : tensor<784x1024xf32>
    %57 = mhlo.add %32, %56 : tensor<784x1024xf32>
    %58 = mhlo.multiply %57, %43 : tensor<784x1024xf32>
    %59 = mhlo.add %31, %58 : tensor<784x1024xf32>
    %60 = mhlo.multiply %59, %arg0 : tensor<784x1024xf32>
    %61 = "mhlo.select"(%24, %25, %60) : (tensor<784x1024xi1>, tensor<784x1024xf32>, tensor<784x1024xf32>) -> tensor<784x1024xf32>
    return %61 : tensor<784x1024xf32>
  }
  func private @jit_clip(%arg0: tensor<784x1024xf32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<784x1024xf32> {
    %0 = "mhlo.broadcast_in_dim"(%arg1) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<784x1024xf32>
    %1 = mhlo.maximum %0, %arg0 : tensor<784x1024xf32>
    %2 = "mhlo.broadcast_in_dim"(%arg2) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<784x1024xf32>
    %3 = mhlo.minimum %2, %1 : tensor<784x1024xf32>
    return %3 : tensor<784x1024xf32>
  }
  func private @main_3(%arg0: tensor<1024xf32>) -> tensor<1024xf32> {
    %0 = mhlo.constant dense<3.000000e+00> : tensor<1024xf32>
    %1 = mhlo.constant dense<2.500000e+00> : tensor<1024xf32>
    %2 = mhlo.constant dense<-2.00214257E-4> : tensor<1024xf32>
    %3 = mhlo.constant dense<2.81022636E-8> : tensor<1024xf32>
    %4 = mhlo.constant dense<1.00950558E-4> : tensor<1024xf32>
    %5 = mhlo.constant dense<3.43273939E-7> : tensor<1024xf32>
    %6 = mhlo.constant dense<0.00134934322> : tensor<1024xf32>
    %7 = mhlo.constant dense<-3.5233877E-6> : tensor<1024xf32>
    %8 = mhlo.constant dense<-0.00367342844> : tensor<1024xf32>
    %9 = mhlo.constant dense<-4.39150654E-6> : tensor<1024xf32>
    %10 = mhlo.constant dense<0.00573950773> : tensor<1024xf32>
    %11 = mhlo.constant dense<2.1858087E-4> : tensor<1024xf32>
    %12 = mhlo.constant dense<-0.0076224613> : tensor<1024xf32>
    %13 = mhlo.constant dense<-0.00125372503> : tensor<1024xf32>
    %14 = mhlo.constant dense<0.00943887047> : tensor<1024xf32>
    %15 = mhlo.constant dense<-0.00417768164> : tensor<1024xf32>
    %16 = mhlo.constant dense<1.00167406> : tensor<1024xf32>
    %17 = mhlo.constant dense<0.246640727> : tensor<1024xf32>
    %18 = mhlo.constant dense<2.83297682> : tensor<1024xf32>
    %19 = mhlo.constant dense<1.50140941> : tensor<1024xf32>
    %20 = mhlo.constant dense<5.000000e+00> : tensor<1024xf32>
    %21 = mhlo.constant dense<0x7F800000> : tensor<1024xf32>
    %22 = mhlo.constant dense<1.000000e+00> : tensor<1024xf32>
    %23 = "mhlo.abs"(%arg0) : (tensor<1024xf32>) -> tensor<1024xf32>
    %24 = "mhlo.compare"(%23, %22) {comparison_direction = "EQ"} : (tensor<1024xf32>, tensor<1024xf32>) -> tensor<1024xi1>
    %25 = mhlo.multiply %arg0, %21 : tensor<1024xf32>
    %26 = "mhlo.negate"(%arg0) : (tensor<1024xf32>) -> tensor<1024xf32>
    %27 = mhlo.multiply %26, %arg0 : tensor<1024xf32>
    %28 = "mhlo.log_plus_one"(%27) : (tensor<1024xf32>) -> tensor<1024xf32>
    %29 = "mhlo.negate"(%28) : (tensor<1024xf32>) -> tensor<1024xf32>
    %30 = "mhlo.compare"(%29, %20) {comparison_direction = "LT"} : (tensor<1024xf32>, tensor<1024xf32>) -> tensor<1024xi1>
    %31 = "mhlo.select"(%30, %19, %18) : (tensor<1024xi1>, tensor<1024xf32>, tensor<1024xf32>) -> tensor<1024xf32>
    %32 = "mhlo.select"(%30, %17, %16) : (tensor<1024xi1>, tensor<1024xf32>, tensor<1024xf32>) -> tensor<1024xf32>
    %33 = "mhlo.select"(%30, %15, %14) : (tensor<1024xi1>, tensor<1024xf32>, tensor<1024xf32>) -> tensor<1024xf32>
    %34 = "mhlo.select"(%30, %13, %12) : (tensor<1024xi1>, tensor<1024xf32>, tensor<1024xf32>) -> tensor<1024xf32>
    %35 = "mhlo.select"(%30, %11, %10) : (tensor<1024xi1>, tensor<1024xf32>, tensor<1024xf32>) -> tensor<1024xf32>
    %36 = "mhlo.select"(%30, %9, %8) : (tensor<1024xi1>, tensor<1024xf32>, tensor<1024xf32>) -> tensor<1024xf32>
    %37 = "mhlo.select"(%30, %7, %6) : (tensor<1024xi1>, tensor<1024xf32>, tensor<1024xf32>) -> tensor<1024xf32>
    %38 = "mhlo.select"(%30, %5, %4) : (tensor<1024xi1>, tensor<1024xf32>, tensor<1024xf32>) -> tensor<1024xf32>
    %39 = "mhlo.select"(%30, %3, %2) : (tensor<1024xi1>, tensor<1024xf32>, tensor<1024xf32>) -> tensor<1024xf32>
    %40 = mhlo.subtract %29, %1 : tensor<1024xf32>
    %41 = "mhlo.sqrt"(%29) : (tensor<1024xf32>) -> tensor<1024xf32>
    %42 = mhlo.subtract %41, %0 : tensor<1024xf32>
    %43 = "mhlo.select"(%30, %40, %42) : (tensor<1024xi1>, tensor<1024xf32>, tensor<1024xf32>) -> tensor<1024xf32>
    %44 = mhlo.multiply %39, %43 : tensor<1024xf32>
    %45 = mhlo.add %38, %44 : tensor<1024xf32>
    %46 = mhlo.multiply %45, %43 : tensor<1024xf32>
    %47 = mhlo.add %37, %46 : tensor<1024xf32>
    %48 = mhlo.multiply %47, %43 : tensor<1024xf32>
    %49 = mhlo.add %36, %48 : tensor<1024xf32>
    %50 = mhlo.multiply %49, %43 : tensor<1024xf32>
    %51 = mhlo.add %35, %50 : tensor<1024xf32>
    %52 = mhlo.multiply %51, %43 : tensor<1024xf32>
    %53 = mhlo.add %34, %52 : tensor<1024xf32>
    %54 = mhlo.multiply %53, %43 : tensor<1024xf32>
    %55 = mhlo.add %33, %54 : tensor<1024xf32>
    %56 = mhlo.multiply %55, %43 : tensor<1024xf32>
    %57 = mhlo.add %32, %56 : tensor<1024xf32>
    %58 = mhlo.multiply %57, %43 : tensor<1024xf32>
    %59 = mhlo.add %31, %58 : tensor<1024xf32>
    %60 = mhlo.multiply %59, %arg0 : tensor<1024xf32>
    %61 = "mhlo.select"(%24, %25, %60) : (tensor<1024xi1>, tensor<1024xf32>, tensor<1024xf32>) -> tensor<1024xf32>
    return %61 : tensor<1024xf32>
  }
  func private @main_4(%arg0: tensor<f32>) -> tensor<f32> {
    %0 = mhlo.constant dense<-0.0142647391> : tensor<f32>
    %1 = mhlo.constant dense<-0.00737332925> : tensor<f32>
    %2 = mhlo.constant dense<-0.00168282702> : tensor<f32>
    %3 = mhlo.constant dense<-2.13374049E-4> : tensor<f32>
    %4 = mhlo.constant dense<-1.45660715E-5> : tensor<f32>
    %5 = mhlo.constant dense<0.000000e+00> : tensor<f32>
    %6 = mhlo.constant dense<-0.0160960332> : tensor<f32>
    %7 = mhlo.constant dense<-2.954600e-03> : tensor<f32>
    %8 = mhlo.constant dense<-7.34990637E-4> : tensor<f32>
    %9 = mhlo.constant dense<-5.69250624E-5> : tensor<f32>
    %10 = mhlo.constant dense<-2.10102394E-6> : tensor<f32>
    %11 = mhlo.constant dense<2.77068146E-8> : tensor<f32>
    %12 = mhlo.constant dense<-2.72614237E-10> : tensor<f32>
    %13 = mhlo.constant dense<4.000000e+00> : tensor<f32>
    %14 = mhlo.constant dense<-4.000000e+00> : tensor<f32>
    %15 = "mhlo.clamp"(%14, %arg0, %13) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
    %16 = mhlo.multiply %15, %15 : tensor<f32>
    %17 = mhlo.multiply %16, %5 : tensor<f32>
    %18 = mhlo.add %17, %12 : tensor<f32>
    %19 = mhlo.multiply %18, %16 : tensor<f32>
    %20 = mhlo.add %19, %11 : tensor<f32>
    %21 = mhlo.multiply %20, %16 : tensor<f32>
    %22 = mhlo.add %21, %10 : tensor<f32>
    %23 = mhlo.multiply %22, %16 : tensor<f32>
    %24 = mhlo.add %23, %9 : tensor<f32>
    %25 = mhlo.multiply %24, %16 : tensor<f32>
    %26 = mhlo.add %25, %8 : tensor<f32>
    %27 = mhlo.multiply %26, %16 : tensor<f32>
    %28 = mhlo.add %27, %7 : tensor<f32>
    %29 = mhlo.multiply %28, %16 : tensor<f32>
    %30 = mhlo.add %29, %6 : tensor<f32>
    %31 = mhlo.multiply %15, %30 : tensor<f32>
    %32 = mhlo.multiply %16, %5 : tensor<f32>
    %33 = mhlo.add %32, %4 : tensor<f32>
    %34 = mhlo.multiply %33, %16 : tensor<f32>
    %35 = mhlo.add %34, %3 : tensor<f32>
    %36 = mhlo.multiply %35, %16 : tensor<f32>
    %37 = mhlo.add %36, %2 : tensor<f32>
    %38 = mhlo.multiply %37, %16 : tensor<f32>
    %39 = mhlo.add %38, %1 : tensor<f32>
    %40 = mhlo.multiply %39, %16 : tensor<f32>
    %41 = mhlo.add %40, %0 : tensor<f32>
    %42 = mhlo.divide %31, %41 : tensor<f32>
    return %42 : tensor<f32>
  }
  func private @main_5(%arg0: tensor<f32>) -> tensor<f32> {
    %0 = mhlo.constant dense<-0.0142647391> : tensor<f32>
    %1 = mhlo.constant dense<-0.00737332925> : tensor<f32>
    %2 = mhlo.constant dense<-0.00168282702> : tensor<f32>
    %3 = mhlo.constant dense<-2.13374049E-4> : tensor<f32>
    %4 = mhlo.constant dense<-1.45660715E-5> : tensor<f32>
    %5 = mhlo.constant dense<0.000000e+00> : tensor<f32>
    %6 = mhlo.constant dense<-0.0160960332> : tensor<f32>
    %7 = mhlo.constant dense<-2.954600e-03> : tensor<f32>
    %8 = mhlo.constant dense<-7.34990637E-4> : tensor<f32>
    %9 = mhlo.constant dense<-5.69250624E-5> : tensor<f32>
    %10 = mhlo.constant dense<-2.10102394E-6> : tensor<f32>
    %11 = mhlo.constant dense<2.77068146E-8> : tensor<f32>
    %12 = mhlo.constant dense<-2.72614237E-10> : tensor<f32>
    %13 = mhlo.constant dense<4.000000e+00> : tensor<f32>
    %14 = mhlo.constant dense<-4.000000e+00> : tensor<f32>
    %15 = "mhlo.clamp"(%14, %arg0, %13) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
    %16 = mhlo.multiply %15, %15 : tensor<f32>
    %17 = mhlo.multiply %16, %5 : tensor<f32>
    %18 = mhlo.add %17, %12 : tensor<f32>
    %19 = mhlo.multiply %18, %16 : tensor<f32>
    %20 = mhlo.add %19, %11 : tensor<f32>
    %21 = mhlo.multiply %20, %16 : tensor<f32>
    %22 = mhlo.add %21, %10 : tensor<f32>
    %23 = mhlo.multiply %22, %16 : tensor<f32>
    %24 = mhlo.add %23, %9 : tensor<f32>
    %25 = mhlo.multiply %24, %16 : tensor<f32>
    %26 = mhlo.add %25, %8 : tensor<f32>
    %27 = mhlo.multiply %26, %16 : tensor<f32>
    %28 = mhlo.add %27, %7 : tensor<f32>
    %29 = mhlo.multiply %28, %16 : tensor<f32>
    %30 = mhlo.add %29, %6 : tensor<f32>
    %31 = mhlo.multiply %15, %30 : tensor<f32>
    %32 = mhlo.multiply %16, %5 : tensor<f32>
    %33 = mhlo.add %32, %4 : tensor<f32>
    %34 = mhlo.multiply %33, %16 : tensor<f32>
    %35 = mhlo.add %34, %3 : tensor<f32>
    %36 = mhlo.multiply %35, %16 : tensor<f32>
    %37 = mhlo.add %36, %2 : tensor<f32>
    %38 = mhlo.multiply %37, %16 : tensor<f32>
    %39 = mhlo.add %38, %1 : tensor<f32>
    %40 = mhlo.multiply %39, %16 : tensor<f32>
    %41 = mhlo.add %40, %0 : tensor<f32>
    %42 = mhlo.divide %31, %41 : tensor<f32>
    return %42 : tensor<f32>
  }
  func private @main_6(%arg0: tensor<1024x1024xf32>) -> tensor<1024x1024xf32> {
    %0 = mhlo.constant dense<3.000000e+00> : tensor<1024x1024xf32>
    %1 = mhlo.constant dense<2.500000e+00> : tensor<1024x1024xf32>
    %2 = mhlo.constant dense<-2.00214257E-4> : tensor<1024x1024xf32>
    %3 = mhlo.constant dense<2.81022636E-8> : tensor<1024x1024xf32>
    %4 = mhlo.constant dense<1.00950558E-4> : tensor<1024x1024xf32>
    %5 = mhlo.constant dense<3.43273939E-7> : tensor<1024x1024xf32>
    %6 = mhlo.constant dense<0.00134934322> : tensor<1024x1024xf32>
    %7 = mhlo.constant dense<-3.5233877E-6> : tensor<1024x1024xf32>
    %8 = mhlo.constant dense<-0.00367342844> : tensor<1024x1024xf32>
    %9 = mhlo.constant dense<-4.39150654E-6> : tensor<1024x1024xf32>
    %10 = mhlo.constant dense<0.00573950773> : tensor<1024x1024xf32>
    %11 = mhlo.constant dense<2.1858087E-4> : tensor<1024x1024xf32>
    %12 = mhlo.constant dense<-0.0076224613> : tensor<1024x1024xf32>
    %13 = mhlo.constant dense<-0.00125372503> : tensor<1024x1024xf32>
    %14 = mhlo.constant dense<0.00943887047> : tensor<1024x1024xf32>
    %15 = mhlo.constant dense<-0.00417768164> : tensor<1024x1024xf32>
    %16 = mhlo.constant dense<1.00167406> : tensor<1024x1024xf32>
    %17 = mhlo.constant dense<0.246640727> : tensor<1024x1024xf32>
    %18 = mhlo.constant dense<2.83297682> : tensor<1024x1024xf32>
    %19 = mhlo.constant dense<1.50140941> : tensor<1024x1024xf32>
    %20 = mhlo.constant dense<5.000000e+00> : tensor<1024x1024xf32>
    %21 = mhlo.constant dense<0x7F800000> : tensor<1024x1024xf32>
    %22 = mhlo.constant dense<1.000000e+00> : tensor<1024x1024xf32>
    %23 = "mhlo.abs"(%arg0) : (tensor<1024x1024xf32>) -> tensor<1024x1024xf32>
    %24 = "mhlo.compare"(%23, %22) {comparison_direction = "EQ"} : (tensor<1024x1024xf32>, tensor<1024x1024xf32>) -> tensor<1024x1024xi1>
    %25 = mhlo.multiply %arg0, %21 : tensor<1024x1024xf32>
    %26 = "mhlo.negate"(%arg0) : (tensor<1024x1024xf32>) -> tensor<1024x1024xf32>
    %27 = mhlo.multiply %26, %arg0 : tensor<1024x1024xf32>
    %28 = "mhlo.log_plus_one"(%27) : (tensor<1024x1024xf32>) -> tensor<1024x1024xf32>
    %29 = "mhlo.negate"(%28) : (tensor<1024x1024xf32>) -> tensor<1024x1024xf32>
    %30 = "mhlo.compare"(%29, %20) {comparison_direction = "LT"} : (tensor<1024x1024xf32>, tensor<1024x1024xf32>) -> tensor<1024x1024xi1>
    %31 = "mhlo.select"(%30, %19, %18) : (tensor<1024x1024xi1>, tensor<1024x1024xf32>, tensor<1024x1024xf32>) -> tensor<1024x1024xf32>
    %32 = "mhlo.select"(%30, %17, %16) : (tensor<1024x1024xi1>, tensor<1024x1024xf32>, tensor<1024x1024xf32>) -> tensor<1024x1024xf32>
    %33 = "mhlo.select"(%30, %15, %14) : (tensor<1024x1024xi1>, tensor<1024x1024xf32>, tensor<1024x1024xf32>) -> tensor<1024x1024xf32>
    %34 = "mhlo.select"(%30, %13, %12) : (tensor<1024x1024xi1>, tensor<1024x1024xf32>, tensor<1024x1024xf32>) -> tensor<1024x1024xf32>
    %35 = "mhlo.select"(%30, %11, %10) : (tensor<1024x1024xi1>, tensor<1024x1024xf32>, tensor<1024x1024xf32>) -> tensor<1024x1024xf32>
    %36 = "mhlo.select"(%30, %9, %8) : (tensor<1024x1024xi1>, tensor<1024x1024xf32>, tensor<1024x1024xf32>) -> tensor<1024x1024xf32>
    %37 = "mhlo.select"(%30, %7, %6) : (tensor<1024x1024xi1>, tensor<1024x1024xf32>, tensor<1024x1024xf32>) -> tensor<1024x1024xf32>
    %38 = "mhlo.select"(%30, %5, %4) : (tensor<1024x1024xi1>, tensor<1024x1024xf32>, tensor<1024x1024xf32>) -> tensor<1024x1024xf32>
    %39 = "mhlo.select"(%30, %3, %2) : (tensor<1024x1024xi1>, tensor<1024x1024xf32>, tensor<1024x1024xf32>) -> tensor<1024x1024xf32>
    %40 = mhlo.subtract %29, %1 : tensor<1024x1024xf32>
    %41 = "mhlo.sqrt"(%29) : (tensor<1024x1024xf32>) -> tensor<1024x1024xf32>
    %42 = mhlo.subtract %41, %0 : tensor<1024x1024xf32>
    %43 = "mhlo.select"(%30, %40, %42) : (tensor<1024x1024xi1>, tensor<1024x1024xf32>, tensor<1024x1024xf32>) -> tensor<1024x1024xf32>
    %44 = mhlo.multiply %39, %43 : tensor<1024x1024xf32>
    %45 = mhlo.add %38, %44 : tensor<1024x1024xf32>
    %46 = mhlo.multiply %45, %43 : tensor<1024x1024xf32>
    %47 = mhlo.add %37, %46 : tensor<1024x1024xf32>
    %48 = mhlo.multiply %47, %43 : tensor<1024x1024xf32>
    %49 = mhlo.add %36, %48 : tensor<1024x1024xf32>
    %50 = mhlo.multiply %49, %43 : tensor<1024x1024xf32>
    %51 = mhlo.add %35, %50 : tensor<1024x1024xf32>
    %52 = mhlo.multiply %51, %43 : tensor<1024x1024xf32>
    %53 = mhlo.add %34, %52 : tensor<1024x1024xf32>
    %54 = mhlo.multiply %53, %43 : tensor<1024x1024xf32>
    %55 = mhlo.add %33, %54 : tensor<1024x1024xf32>
    %56 = mhlo.multiply %55, %43 : tensor<1024x1024xf32>
    %57 = mhlo.add %32, %56 : tensor<1024x1024xf32>
    %58 = mhlo.multiply %57, %43 : tensor<1024x1024xf32>
    %59 = mhlo.add %31, %58 : tensor<1024x1024xf32>
    %60 = mhlo.multiply %59, %arg0 : tensor<1024x1024xf32>
    %61 = "mhlo.select"(%24, %25, %60) : (tensor<1024x1024xi1>, tensor<1024x1024xf32>, tensor<1024x1024xf32>) -> tensor<1024x1024xf32>
    return %61 : tensor<1024x1024xf32>
  }
  func private @jit_clip_7(%arg0: tensor<1024x1024xf32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<1024x1024xf32> {
    %0 = "mhlo.broadcast_in_dim"(%arg1) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1024x1024xf32>
    %1 = mhlo.maximum %0, %arg0 : tensor<1024x1024xf32>
    %2 = "mhlo.broadcast_in_dim"(%arg2) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1024x1024xf32>
    %3 = mhlo.minimum %2, %1 : tensor<1024x1024xf32>
    return %3 : tensor<1024x1024xf32>
  }
  func private @main_8(%arg0: tensor<1024xf32>) -> tensor<1024xf32> {
    %0 = mhlo.constant dense<3.000000e+00> : tensor<1024xf32>
    %1 = mhlo.constant dense<2.500000e+00> : tensor<1024xf32>
    %2 = mhlo.constant dense<-2.00214257E-4> : tensor<1024xf32>
    %3 = mhlo.constant dense<2.81022636E-8> : tensor<1024xf32>
    %4 = mhlo.constant dense<1.00950558E-4> : tensor<1024xf32>
    %5 = mhlo.constant dense<3.43273939E-7> : tensor<1024xf32>
    %6 = mhlo.constant dense<0.00134934322> : tensor<1024xf32>
    %7 = mhlo.constant dense<-3.5233877E-6> : tensor<1024xf32>
    %8 = mhlo.constant dense<-0.00367342844> : tensor<1024xf32>
    %9 = mhlo.constant dense<-4.39150654E-6> : tensor<1024xf32>
    %10 = mhlo.constant dense<0.00573950773> : tensor<1024xf32>
    %11 = mhlo.constant dense<2.1858087E-4> : tensor<1024xf32>
    %12 = mhlo.constant dense<-0.0076224613> : tensor<1024xf32>
    %13 = mhlo.constant dense<-0.00125372503> : tensor<1024xf32>
    %14 = mhlo.constant dense<0.00943887047> : tensor<1024xf32>
    %15 = mhlo.constant dense<-0.00417768164> : tensor<1024xf32>
    %16 = mhlo.constant dense<1.00167406> : tensor<1024xf32>
    %17 = mhlo.constant dense<0.246640727> : tensor<1024xf32>
    %18 = mhlo.constant dense<2.83297682> : tensor<1024xf32>
    %19 = mhlo.constant dense<1.50140941> : tensor<1024xf32>
    %20 = mhlo.constant dense<5.000000e+00> : tensor<1024xf32>
    %21 = mhlo.constant dense<0x7F800000> : tensor<1024xf32>
    %22 = mhlo.constant dense<1.000000e+00> : tensor<1024xf32>
    %23 = "mhlo.abs"(%arg0) : (tensor<1024xf32>) -> tensor<1024xf32>
    %24 = "mhlo.compare"(%23, %22) {comparison_direction = "EQ"} : (tensor<1024xf32>, tensor<1024xf32>) -> tensor<1024xi1>
    %25 = mhlo.multiply %arg0, %21 : tensor<1024xf32>
    %26 = "mhlo.negate"(%arg0) : (tensor<1024xf32>) -> tensor<1024xf32>
    %27 = mhlo.multiply %26, %arg0 : tensor<1024xf32>
    %28 = "mhlo.log_plus_one"(%27) : (tensor<1024xf32>) -> tensor<1024xf32>
    %29 = "mhlo.negate"(%28) : (tensor<1024xf32>) -> tensor<1024xf32>
    %30 = "mhlo.compare"(%29, %20) {comparison_direction = "LT"} : (tensor<1024xf32>, tensor<1024xf32>) -> tensor<1024xi1>
    %31 = "mhlo.select"(%30, %19, %18) : (tensor<1024xi1>, tensor<1024xf32>, tensor<1024xf32>) -> tensor<1024xf32>
    %32 = "mhlo.select"(%30, %17, %16) : (tensor<1024xi1>, tensor<1024xf32>, tensor<1024xf32>) -> tensor<1024xf32>
    %33 = "mhlo.select"(%30, %15, %14) : (tensor<1024xi1>, tensor<1024xf32>, tensor<1024xf32>) -> tensor<1024xf32>
    %34 = "mhlo.select"(%30, %13, %12) : (tensor<1024xi1>, tensor<1024xf32>, tensor<1024xf32>) -> tensor<1024xf32>
    %35 = "mhlo.select"(%30, %11, %10) : (tensor<1024xi1>, tensor<1024xf32>, tensor<1024xf32>) -> tensor<1024xf32>
    %36 = "mhlo.select"(%30, %9, %8) : (tensor<1024xi1>, tensor<1024xf32>, tensor<1024xf32>) -> tensor<1024xf32>
    %37 = "mhlo.select"(%30, %7, %6) : (tensor<1024xi1>, tensor<1024xf32>, tensor<1024xf32>) -> tensor<1024xf32>
    %38 = "mhlo.select"(%30, %5, %4) : (tensor<1024xi1>, tensor<1024xf32>, tensor<1024xf32>) -> tensor<1024xf32>
    %39 = "mhlo.select"(%30, %3, %2) : (tensor<1024xi1>, tensor<1024xf32>, tensor<1024xf32>) -> tensor<1024xf32>
    %40 = mhlo.subtract %29, %1 : tensor<1024xf32>
    %41 = "mhlo.sqrt"(%29) : (tensor<1024xf32>) -> tensor<1024xf32>
    %42 = mhlo.subtract %41, %0 : tensor<1024xf32>
    %43 = "mhlo.select"(%30, %40, %42) : (tensor<1024xi1>, tensor<1024xf32>, tensor<1024xf32>) -> tensor<1024xf32>
    %44 = mhlo.multiply %39, %43 : tensor<1024xf32>
    %45 = mhlo.add %38, %44 : tensor<1024xf32>
    %46 = mhlo.multiply %45, %43 : tensor<1024xf32>
    %47 = mhlo.add %37, %46 : tensor<1024xf32>
    %48 = mhlo.multiply %47, %43 : tensor<1024xf32>
    %49 = mhlo.add %36, %48 : tensor<1024xf32>
    %50 = mhlo.multiply %49, %43 : tensor<1024xf32>
    %51 = mhlo.add %35, %50 : tensor<1024xf32>
    %52 = mhlo.multiply %51, %43 : tensor<1024xf32>
    %53 = mhlo.add %34, %52 : tensor<1024xf32>
    %54 = mhlo.multiply %53, %43 : tensor<1024xf32>
    %55 = mhlo.add %33, %54 : tensor<1024xf32>
    %56 = mhlo.multiply %55, %43 : tensor<1024xf32>
    %57 = mhlo.add %32, %56 : tensor<1024xf32>
    %58 = mhlo.multiply %57, %43 : tensor<1024xf32>
    %59 = mhlo.add %31, %58 : tensor<1024xf32>
    %60 = mhlo.multiply %59, %arg0 : tensor<1024xf32>
    %61 = "mhlo.select"(%24, %25, %60) : (tensor<1024xi1>, tensor<1024xf32>, tensor<1024xf32>) -> tensor<1024xf32>
    return %61 : tensor<1024xf32>
  }
  func private @main_9(%arg0: tensor<f32>) -> tensor<f32> {
    %0 = mhlo.constant dense<-0.0142647391> : tensor<f32>
    %1 = mhlo.constant dense<-0.00737332925> : tensor<f32>
    %2 = mhlo.constant dense<-0.00168282702> : tensor<f32>
    %3 = mhlo.constant dense<-2.13374049E-4> : tensor<f32>
    %4 = mhlo.constant dense<-1.45660715E-5> : tensor<f32>
    %5 = mhlo.constant dense<0.000000e+00> : tensor<f32>
    %6 = mhlo.constant dense<-0.0160960332> : tensor<f32>
    %7 = mhlo.constant dense<-2.954600e-03> : tensor<f32>
    %8 = mhlo.constant dense<-7.34990637E-4> : tensor<f32>
    %9 = mhlo.constant dense<-5.69250624E-5> : tensor<f32>
    %10 = mhlo.constant dense<-2.10102394E-6> : tensor<f32>
    %11 = mhlo.constant dense<2.77068146E-8> : tensor<f32>
    %12 = mhlo.constant dense<-2.72614237E-10> : tensor<f32>
    %13 = mhlo.constant dense<4.000000e+00> : tensor<f32>
    %14 = mhlo.constant dense<-4.000000e+00> : tensor<f32>
    %15 = "mhlo.clamp"(%14, %arg0, %13) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
    %16 = mhlo.multiply %15, %15 : tensor<f32>
    %17 = mhlo.multiply %16, %5 : tensor<f32>
    %18 = mhlo.add %17, %12 : tensor<f32>
    %19 = mhlo.multiply %18, %16 : tensor<f32>
    %20 = mhlo.add %19, %11 : tensor<f32>
    %21 = mhlo.multiply %20, %16 : tensor<f32>
    %22 = mhlo.add %21, %10 : tensor<f32>
    %23 = mhlo.multiply %22, %16 : tensor<f32>
    %24 = mhlo.add %23, %9 : tensor<f32>
    %25 = mhlo.multiply %24, %16 : tensor<f32>
    %26 = mhlo.add %25, %8 : tensor<f32>
    %27 = mhlo.multiply %26, %16 : tensor<f32>
    %28 = mhlo.add %27, %7 : tensor<f32>
    %29 = mhlo.multiply %28, %16 : tensor<f32>
    %30 = mhlo.add %29, %6 : tensor<f32>
    %31 = mhlo.multiply %15, %30 : tensor<f32>
    %32 = mhlo.multiply %16, %5 : tensor<f32>
    %33 = mhlo.add %32, %4 : tensor<f32>
    %34 = mhlo.multiply %33, %16 : tensor<f32>
    %35 = mhlo.add %34, %3 : tensor<f32>
    %36 = mhlo.multiply %35, %16 : tensor<f32>
    %37 = mhlo.add %36, %2 : tensor<f32>
    %38 = mhlo.multiply %37, %16 : tensor<f32>
    %39 = mhlo.add %38, %1 : tensor<f32>
    %40 = mhlo.multiply %39, %16 : tensor<f32>
    %41 = mhlo.add %40, %0 : tensor<f32>
    %42 = mhlo.divide %31, %41 : tensor<f32>
    return %42 : tensor<f32>
  }
  func private @main_10(%arg0: tensor<f32>) -> tensor<f32> {
    %0 = mhlo.constant dense<-0.0142647391> : tensor<f32>
    %1 = mhlo.constant dense<-0.00737332925> : tensor<f32>
    %2 = mhlo.constant dense<-0.00168282702> : tensor<f32>
    %3 = mhlo.constant dense<-2.13374049E-4> : tensor<f32>
    %4 = mhlo.constant dense<-1.45660715E-5> : tensor<f32>
    %5 = mhlo.constant dense<0.000000e+00> : tensor<f32>
    %6 = mhlo.constant dense<-0.0160960332> : tensor<f32>
    %7 = mhlo.constant dense<-2.954600e-03> : tensor<f32>
    %8 = mhlo.constant dense<-7.34990637E-4> : tensor<f32>
    %9 = mhlo.constant dense<-5.69250624E-5> : tensor<f32>
    %10 = mhlo.constant dense<-2.10102394E-6> : tensor<f32>
    %11 = mhlo.constant dense<2.77068146E-8> : tensor<f32>
    %12 = mhlo.constant dense<-2.72614237E-10> : tensor<f32>
    %13 = mhlo.constant dense<4.000000e+00> : tensor<f32>
    %14 = mhlo.constant dense<-4.000000e+00> : tensor<f32>
    %15 = "mhlo.clamp"(%14, %arg0, %13) : (tensor<f32>, tensor<f32>, tensor<f32>) -> tensor<f32>
    %16 = mhlo.multiply %15, %15 : tensor<f32>
    %17 = mhlo.multiply %16, %5 : tensor<f32>
    %18 = mhlo.add %17, %12 : tensor<f32>
    %19 = mhlo.multiply %18, %16 : tensor<f32>
    %20 = mhlo.add %19, %11 : tensor<f32>
    %21 = mhlo.multiply %20, %16 : tensor<f32>
    %22 = mhlo.add %21, %10 : tensor<f32>
    %23 = mhlo.multiply %22, %16 : tensor<f32>
    %24 = mhlo.add %23, %9 : tensor<f32>
    %25 = mhlo.multiply %24, %16 : tensor<f32>
    %26 = mhlo.add %25, %8 : tensor<f32>
    %27 = mhlo.multiply %26, %16 : tensor<f32>
    %28 = mhlo.add %27, %7 : tensor<f32>
    %29 = mhlo.multiply %28, %16 : tensor<f32>
    %30 = mhlo.add %29, %6 : tensor<f32>
    %31 = mhlo.multiply %15, %30 : tensor<f32>
    %32 = mhlo.multiply %16, %5 : tensor<f32>
    %33 = mhlo.add %32, %4 : tensor<f32>
    %34 = mhlo.multiply %33, %16 : tensor<f32>
    %35 = mhlo.add %34, %3 : tensor<f32>
    %36 = mhlo.multiply %35, %16 : tensor<f32>
    %37 = mhlo.add %36, %2 : tensor<f32>
    %38 = mhlo.multiply %37, %16 : tensor<f32>
    %39 = mhlo.add %38, %1 : tensor<f32>
    %40 = mhlo.multiply %39, %16 : tensor<f32>
    %41 = mhlo.add %40, %0 : tensor<f32>
    %42 = mhlo.divide %31, %41 : tensor<f32>
    return %42 : tensor<f32>
  }
  func private @main_11(%arg0: tensor<1024x10xf32>) -> tensor<1024x10xf32> {
    %0 = mhlo.constant dense<3.000000e+00> : tensor<1024x10xf32>
    %1 = mhlo.constant dense<2.500000e+00> : tensor<1024x10xf32>
    %2 = mhlo.constant dense<-2.00214257E-4> : tensor<1024x10xf32>
    %3 = mhlo.constant dense<2.81022636E-8> : tensor<1024x10xf32>
    %4 = mhlo.constant dense<1.00950558E-4> : tensor<1024x10xf32>
    %5 = mhlo.constant dense<3.43273939E-7> : tensor<1024x10xf32>
    %6 = mhlo.constant dense<0.00134934322> : tensor<1024x10xf32>
    %7 = mhlo.constant dense<-3.5233877E-6> : tensor<1024x10xf32>
    %8 = mhlo.constant dense<-0.00367342844> : tensor<1024x10xf32>
    %9 = mhlo.constant dense<-4.39150654E-6> : tensor<1024x10xf32>
    %10 = mhlo.constant dense<0.00573950773> : tensor<1024x10xf32>
    %11 = mhlo.constant dense<2.1858087E-4> : tensor<1024x10xf32>
    %12 = mhlo.constant dense<-0.0076224613> : tensor<1024x10xf32>
    %13 = mhlo.constant dense<-0.00125372503> : tensor<1024x10xf32>
    %14 = mhlo.constant dense<0.00943887047> : tensor<1024x10xf32>
    %15 = mhlo.constant dense<-0.00417768164> : tensor<1024x10xf32>
    %16 = mhlo.constant dense<1.00167406> : tensor<1024x10xf32>
    %17 = mhlo.constant dense<0.246640727> : tensor<1024x10xf32>
    %18 = mhlo.constant dense<2.83297682> : tensor<1024x10xf32>
    %19 = mhlo.constant dense<1.50140941> : tensor<1024x10xf32>
    %20 = mhlo.constant dense<5.000000e+00> : tensor<1024x10xf32>
    %21 = mhlo.constant dense<0x7F800000> : tensor<1024x10xf32>
    %22 = mhlo.constant dense<1.000000e+00> : tensor<1024x10xf32>
    %23 = "mhlo.abs"(%arg0) : (tensor<1024x10xf32>) -> tensor<1024x10xf32>
    %24 = "mhlo.compare"(%23, %22) {comparison_direction = "EQ"} : (tensor<1024x10xf32>, tensor<1024x10xf32>) -> tensor<1024x10xi1>
    %25 = mhlo.multiply %arg0, %21 : tensor<1024x10xf32>
    %26 = "mhlo.negate"(%arg0) : (tensor<1024x10xf32>) -> tensor<1024x10xf32>
    %27 = mhlo.multiply %26, %arg0 : tensor<1024x10xf32>
    %28 = "mhlo.log_plus_one"(%27) : (tensor<1024x10xf32>) -> tensor<1024x10xf32>
    %29 = "mhlo.negate"(%28) : (tensor<1024x10xf32>) -> tensor<1024x10xf32>
    %30 = "mhlo.compare"(%29, %20) {comparison_direction = "LT"} : (tensor<1024x10xf32>, tensor<1024x10xf32>) -> tensor<1024x10xi1>
    %31 = "mhlo.select"(%30, %19, %18) : (tensor<1024x10xi1>, tensor<1024x10xf32>, tensor<1024x10xf32>) -> tensor<1024x10xf32>
    %32 = "mhlo.select"(%30, %17, %16) : (tensor<1024x10xi1>, tensor<1024x10xf32>, tensor<1024x10xf32>) -> tensor<1024x10xf32>
    %33 = "mhlo.select"(%30, %15, %14) : (tensor<1024x10xi1>, tensor<1024x10xf32>, tensor<1024x10xf32>) -> tensor<1024x10xf32>
    %34 = "mhlo.select"(%30, %13, %12) : (tensor<1024x10xi1>, tensor<1024x10xf32>, tensor<1024x10xf32>) -> tensor<1024x10xf32>
    %35 = "mhlo.select"(%30, %11, %10) : (tensor<1024x10xi1>, tensor<1024x10xf32>, tensor<1024x10xf32>) -> tensor<1024x10xf32>
    %36 = "mhlo.select"(%30, %9, %8) : (tensor<1024x10xi1>, tensor<1024x10xf32>, tensor<1024x10xf32>) -> tensor<1024x10xf32>
    %37 = "mhlo.select"(%30, %7, %6) : (tensor<1024x10xi1>, tensor<1024x10xf32>, tensor<1024x10xf32>) -> tensor<1024x10xf32>
    %38 = "mhlo.select"(%30, %5, %4) : (tensor<1024x10xi1>, tensor<1024x10xf32>, tensor<1024x10xf32>) -> tensor<1024x10xf32>
    %39 = "mhlo.select"(%30, %3, %2) : (tensor<1024x10xi1>, tensor<1024x10xf32>, tensor<1024x10xf32>) -> tensor<1024x10xf32>
    %40 = mhlo.subtract %29, %1 : tensor<1024x10xf32>
    %41 = "mhlo.sqrt"(%29) : (tensor<1024x10xf32>) -> tensor<1024x10xf32>
    %42 = mhlo.subtract %41, %0 : tensor<1024x10xf32>
    %43 = "mhlo.select"(%30, %40, %42) : (tensor<1024x10xi1>, tensor<1024x10xf32>, tensor<1024x10xf32>) -> tensor<1024x10xf32>
    %44 = mhlo.multiply %39, %43 : tensor<1024x10xf32>
    %45 = mhlo.add %38, %44 : tensor<1024x10xf32>
    %46 = mhlo.multiply %45, %43 : tensor<1024x10xf32>
    %47 = mhlo.add %37, %46 : tensor<1024x10xf32>
    %48 = mhlo.multiply %47, %43 : tensor<1024x10xf32>
    %49 = mhlo.add %36, %48 : tensor<1024x10xf32>
    %50 = mhlo.multiply %49, %43 : tensor<1024x10xf32>
    %51 = mhlo.add %35, %50 : tensor<1024x10xf32>
    %52 = mhlo.multiply %51, %43 : tensor<1024x10xf32>
    %53 = mhlo.add %34, %52 : tensor<1024x10xf32>
    %54 = mhlo.multiply %53, %43 : tensor<1024x10xf32>
    %55 = mhlo.add %33, %54 : tensor<1024x10xf32>
    %56 = mhlo.multiply %55, %43 : tensor<1024x10xf32>
    %57 = mhlo.add %32, %56 : tensor<1024x10xf32>
    %58 = mhlo.multiply %57, %43 : tensor<1024x10xf32>
    %59 = mhlo.add %31, %58 : tensor<1024x10xf32>
    %60 = mhlo.multiply %59, %arg0 : tensor<1024x10xf32>
    %61 = "mhlo.select"(%24, %25, %60) : (tensor<1024x10xi1>, tensor<1024x10xf32>, tensor<1024x10xf32>) -> tensor<1024x10xf32>
    return %61 : tensor<1024x10xf32>
  }
  func private @jit_clip_12(%arg0: tensor<1024x10xf32>, %arg1: tensor<f32>, %arg2: tensor<f32>) -> tensor<1024x10xf32> {
    %0 = "mhlo.broadcast_in_dim"(%arg1) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1024x10xf32>
    %1 = mhlo.maximum %0, %arg0 : tensor<1024x10xf32>
    %2 = "mhlo.broadcast_in_dim"(%arg2) {broadcast_dimensions = dense<> : tensor<0xi64>} : (tensor<f32>) -> tensor<1024x10xf32>
    %3 = mhlo.minimum %2, %1 : tensor<1024x10xf32>
    return %3 : tensor<1024x10xf32>
  }
  func private @main_13(%arg0: tensor<10xf32>) -> tensor<10xf32> {
    %0 = mhlo.constant dense<3.000000e+00> : tensor<10xf32>
    %1 = mhlo.constant dense<2.500000e+00> : tensor<10xf32>
    %2 = mhlo.constant dense<-2.00214257E-4> : tensor<10xf32>
    %3 = mhlo.constant dense<2.81022636E-8> : tensor<10xf32>
    %4 = mhlo.constant dense<1.00950558E-4> : tensor<10xf32>
    %5 = mhlo.constant dense<3.43273939E-7> : tensor<10xf32>
    %6 = mhlo.constant dense<0.00134934322> : tensor<10xf32>
    %7 = mhlo.constant dense<-3.5233877E-6> : tensor<10xf32>
    %8 = mhlo.constant dense<-0.00367342844> : tensor<10xf32>
    %9 = mhlo.constant dense<-4.39150654E-6> : tensor<10xf32>
    %10 = mhlo.constant dense<0.00573950773> : tensor<10xf32>
    %11 = mhlo.constant dense<2.1858087E-4> : tensor<10xf32>
    %12 = mhlo.constant dense<-0.0076224613> : tensor<10xf32>
    %13 = mhlo.constant dense<-0.00125372503> : tensor<10xf32>
    %14 = mhlo.constant dense<0.00943887047> : tensor<10xf32>
    %15 = mhlo.constant dense<-0.00417768164> : tensor<10xf32>
    %16 = mhlo.constant dense<1.00167406> : tensor<10xf32>
    %17 = mhlo.constant dense<0.246640727> : tensor<10xf32>
    %18 = mhlo.constant dense<2.83297682> : tensor<10xf32>
    %19 = mhlo.constant dense<1.50140941> : tensor<10xf32>
    %20 = mhlo.constant dense<5.000000e+00> : tensor<10xf32>
    %21 = mhlo.constant dense<0x7F800000> : tensor<10xf32>
    %22 = mhlo.constant dense<1.000000e+00> : tensor<10xf32>
    %23 = "mhlo.abs"(%arg0) : (tensor<10xf32>) -> tensor<10xf32>
    %24 = "mhlo.compare"(%23, %22) {comparison_direction = "EQ"} : (tensor<10xf32>, tensor<10xf32>) -> tensor<10xi1>
    %25 = mhlo.multiply %arg0, %21 : tensor<10xf32>
    %26 = "mhlo.negate"(%arg0) : (tensor<10xf32>) -> tensor<10xf32>
    %27 = mhlo.multiply %26, %arg0 : tensor<10xf32>
    %28 = "mhlo.log_plus_one"(%27) : (tensor<10xf32>) -> tensor<10xf32>
    %29 = "mhlo.negate"(%28) : (tensor<10xf32>) -> tensor<10xf32>
    %30 = "mhlo.compare"(%29, %20) {comparison_direction = "LT"} : (tensor<10xf32>, tensor<10xf32>) -> tensor<10xi1>
    %31 = "mhlo.select"(%30, %19, %18) : (tensor<10xi1>, tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
    %32 = "mhlo.select"(%30, %17, %16) : (tensor<10xi1>, tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
    %33 = "mhlo.select"(%30, %15, %14) : (tensor<10xi1>, tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
    %34 = "mhlo.select"(%30, %13, %12) : (tensor<10xi1>, tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
    %35 = "mhlo.select"(%30, %11, %10) : (tensor<10xi1>, tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
    %36 = "mhlo.select"(%30, %9, %8) : (tensor<10xi1>, tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
    %37 = "mhlo.select"(%30, %7, %6) : (tensor<10xi1>, tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
    %38 = "mhlo.select"(%30, %5, %4) : (tensor<10xi1>, tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
    %39 = "mhlo.select"(%30, %3, %2) : (tensor<10xi1>, tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
    %40 = mhlo.subtract %29, %1 : tensor<10xf32>
    %41 = "mhlo.sqrt"(%29) : (tensor<10xf32>) -> tensor<10xf32>
    %42 = mhlo.subtract %41, %0 : tensor<10xf32>
    %43 = "mhlo.select"(%30, %40, %42) : (tensor<10xi1>, tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
    %44 = mhlo.multiply %39, %43 : tensor<10xf32>
    %45 = mhlo.add %38, %44 : tensor<10xf32>
    %46 = mhlo.multiply %45, %43 : tensor<10xf32>
    %47 = mhlo.add %37, %46 : tensor<10xf32>
    %48 = mhlo.multiply %47, %43 : tensor<10xf32>
    %49 = mhlo.add %36, %48 : tensor<10xf32>
    %50 = mhlo.multiply %49, %43 : tensor<10xf32>
    %51 = mhlo.add %35, %50 : tensor<10xf32>
    %52 = mhlo.multiply %51, %43 : tensor<10xf32>
    %53 = mhlo.add %34, %52 : tensor<10xf32>
    %54 = mhlo.multiply %53, %43 : tensor<10xf32>
    %55 = mhlo.add %33, %54 : tensor<10xf32>
    %56 = mhlo.multiply %55, %43 : tensor<10xf32>
    %57 = mhlo.add %32, %56 : tensor<10xf32>
    %58 = mhlo.multiply %57, %43 : tensor<10xf32>
    %59 = mhlo.add %31, %58 : tensor<10xf32>
    %60 = mhlo.multiply %59, %arg0 : tensor<10xf32>
    %61 = "mhlo.select"(%24, %25, %60) : (tensor<10xi1>, tensor<10xf32>, tensor<10xf32>) -> tensor<10xf32>
    return %61 : tensor<10xf32>
  }
  func @update(%arg0: tensor<128x784xf32>, %arg1: tensor<128x10xf32>) {
    %0 = iree_input.global.load @init_params$0 : tensor<784x1024xf32>
    %1 = iree_input.global.load @opt_state$1 : tensor<784x1024xf32>
    %2 = iree_input.global.load @init_params$1 : tensor<1024xf32>
    %3 = iree_input.global.load @opt_state$3 : tensor<1024xf32>
    %4 = iree_input.global.load @init_params$2 : tensor<1024x1024xf32>
    %5 = iree_input.global.load @opt_state$5 : tensor<1024x1024xf32>
    %6 = iree_input.global.load @init_params$3 : tensor<1024xf32>
    %7 = iree_input.global.load @opt_state$7 : tensor<1024xf32>
    %8 = iree_input.global.load @init_params$4 : tensor<1024x10xf32>
    %9 = iree_input.global.load @opt_state$9 : tensor<1024x10xf32>
    %10 = iree_input.global.load @init_params$5 : tensor<10xf32>
    %11 = iree_input.global.load @opt_state$11 : tensor<10xf32>
    %12:12 = call @main_7(%arg0, %arg1, %0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11) : (tensor<128x784xf32>, tensor<128x10xf32>, tensor<784x1024xf32>, tensor<784x1024xf32>, tensor<1024xf32>, tensor<1024xf32>, tensor<1024x1024xf32>, tensor<1024x1024xf32>, tensor<1024xf32>, tensor<1024xf32>, tensor<1024x10xf32>, tensor<1024x10xf32>, tensor<10xf32>, tensor<10xf32>) -> (tensor<784x1024xf32>, tensor<784x1024xf32>, tensor<1024xf32>, tensor<1024xf32>, tensor<1024x1024xf32>, tensor<1024x1024xf32>, tensor<1024xf32>, tensor<1024xf32>, tensor<1024x10xf32>, tensor<1024x10xf32>, tensor<10xf32>, tensor<10xf32>)
    iree_input.global.store %12#0, @init_params$0 : tensor<784x1024xf32>
    iree_input.global.store %12#1, @opt_state$1 : tensor<784x1024xf32>
    iree_input.global.store %12#2, @init_params$1 : tensor<1024xf32>
    iree_input.global.store %12#3, @opt_state$3 : tensor<1024xf32>
    iree_input.global.store %12#4, @init_params$2 : tensor<1024x1024xf32>
    iree_input.global.store %12#5, @opt_state$5 : tensor<1024x1024xf32>
    iree_input.global.store %12#6, @init_params$3 : tensor<1024xf32>
    iree_input.global.store %12#7, @opt_state$7 : tensor<1024xf32>
    iree_input.global.store %12#8, @init_params$4 : tensor<1024x10xf32>
    iree_input.global.store %12#9, @opt_state$9 : tensor<1024x10xf32>
    iree_input.global.store %12#10, @init_params$5 : tensor<10xf32>
    iree_input.global.store %12#11, @opt_state$11 : tensor<10xf32>
    return
  }
  func private @main_7(%arg0: tensor<128x784xf32>, %arg1: tensor<128x10xf32>, %arg2: tensor<784x1024xf32>, %arg3: tensor<784x1024xf32>, %arg4: tensor<1024xf32>, %arg5: tensor<1024xf32>, %arg6: tensor<1024x1024xf32>, %arg7: tensor<1024x1024xf32>, %arg8: tensor<1024xf32>, %arg9: tensor<1024xf32>, %arg10: tensor<1024x10xf32>, %arg11: tensor<1024x10xf32>, %arg12: tensor<10xf32>, %arg13: tensor<10xf32>) -> (tensor<784x1024xf32>, tensor<784x1024xf32>, tensor<1024xf32>, tensor<1024xf32>, tensor<1024x1024xf32>, tensor<1024x1024xf32>, tensor<1024xf32>, tensor<1024xf32>, tensor<1024x10xf32>, tensor<1024x10xf32>, tensor<10xf32>, tensor<10xf32>) {
    %0 = mhlo.constant dense<1.000000e-03> : tensor<10xf32>
    %1 = mhlo.constant dense<0.899999976> : tensor<10xf32>
    %2 = mhlo.constant dense<1.000000e-03> : tensor<1024x10xf32>
    %3 = mhlo.constant dense<0.899999976> : tensor<1024x10xf32>
    %4 = mhlo.constant dense<1.000000e-03> : tensor<1024xf32>
    %5 = mhlo.constant dense<0.899999976> : tensor<1024xf32>
    %6 = mhlo.constant dense<1.000000e-03> : tensor<1024x1024xf32>
    %7 = mhlo.constant dense<0.899999976> : tensor<1024x1024xf32>
    %8 = mhlo.constant dense<1.000000e-03> : tensor<784x1024xf32>
    %9 = mhlo.constant dense<0.899999976> : tensor<784x1024xf32>
    %10 = mhlo.constant dense<0.000000e+00> : tensor<f32>
    %11 = mhlo.constant dense<0.000000e+00> : tensor<128x1024xf32>
    %12 = mhlo.constant dense<-7.812500e-03> : tensor<128x10xf32>
    %13 = mhlo.constant dense<0xFF800000> : tensor<f32>
    %14 = "mhlo.dot_general"(%arg0, %arg2) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = ["DEFAULT", "DEFAULT"]} : (tensor<128x784xf32>, tensor<784x1024xf32>) -> tensor<128x1024xf32>
    %15 = "mhlo.broadcast_in_dim"(%arg4) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<1024xf32>) -> tensor<1x1024xf32>
    %16 = "mhlo.broadcast_in_dim"(%15) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor<1x1024xf32>) -> tensor<128x1024xf32>
    %17 = mhlo.add %14, %16 : tensor<128x1024xf32>
    %18 = mhlo.maximum %17, %11 : tensor<128x1024xf32>
    %19 = "mhlo.compare"(%17, %11) {compare_type = "FLOAT", comparison_direction = "GT"} : (tensor<128x1024xf32>, tensor<128x1024xf32>) -> tensor<128x1024xi1>
    %20 = "mhlo.dot_general"(%18, %arg6) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = ["DEFAULT", "DEFAULT"]} : (tensor<128x1024xf32>, tensor<1024x1024xf32>) -> tensor<128x1024xf32>
    %21 = "mhlo.broadcast_in_dim"(%arg8) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<1024xf32>) -> tensor<1x1024xf32>
    %22 = "mhlo.broadcast_in_dim"(%21) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor<1x1024xf32>) -> tensor<128x1024xf32>
    %23 = mhlo.add %20, %22 : tensor<128x1024xf32>
    %24 = mhlo.maximum %23, %11 : tensor<128x1024xf32>
    %25 = "mhlo.compare"(%23, %11) {compare_type = "FLOAT", comparison_direction = "GT"} : (tensor<128x1024xf32>, tensor<128x1024xf32>) -> tensor<128x1024xi1>
    %26 = "mhlo.dot_general"(%24, %arg10) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = ["DEFAULT", "DEFAULT"]} : (tensor<128x1024xf32>, tensor<1024x10xf32>) -> tensor<128x10xf32>
    %27 = "mhlo.broadcast_in_dim"(%arg12) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<10xf32>) -> tensor<1x10xf32>
    %28 = "mhlo.broadcast_in_dim"(%27) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor<1x10xf32>) -> tensor<128x10xf32>
    %29 = mhlo.add %26, %28 : tensor<128x10xf32>
    %30 = "mhlo.reduce"(%29, %13) ( {
    ^bb0(%arg14: tensor<f32>, %arg15: tensor<f32>):  // no predecessors
      %89 = mhlo.maximum %arg14, %arg15 : tensor<f32>
      "mhlo.return"(%89) : (tensor<f32>) -> ()
    }) {dimensions = dense<1> : tensor<1xi64>} : (tensor<128x10xf32>, tensor<f32>) -> tensor<128xf32>
    %31 = "mhlo.broadcast_in_dim"(%30) {broadcast_dimensions = dense<0> : tensor<1xi64>} : (tensor<128xf32>) -> tensor<128x1xf32>
    %32 = "mhlo.broadcast_in_dim"(%31) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor<128x1xf32>) -> tensor<128x10xf32>
    %33 = mhlo.subtract %29, %32 : tensor<128x10xf32>
    %34 = "mhlo.exponential"(%33) : (tensor<128x10xf32>) -> tensor<128x10xf32>
    %35 = "mhlo.reduce"(%34, %10) ( {
    ^bb0(%arg14: tensor<f32>, %arg15: tensor<f32>):  // no predecessors
      %89 = mhlo.add %arg14, %arg15 : tensor<f32>
      "mhlo.return"(%89) : (tensor<f32>) -> ()
    }) {dimensions = dense<1> : tensor<1xi64>} : (tensor<128x10xf32>, tensor<f32>) -> tensor<128xf32>
    %36 = "mhlo.broadcast_in_dim"(%35) {broadcast_dimensions = dense<0> : tensor<1xi64>} : (tensor<128xf32>) -> tensor<128x1xf32>
    %37 = mhlo.multiply %arg1, %12 : tensor<128x10xf32>
    %38 = "mhlo.negate"(%37) : (tensor<128x10xf32>) -> tensor<128x10xf32>
    %39 = "mhlo.reduce"(%38, %10) ( {
    ^bb0(%arg14: tensor<f32>, %arg15: tensor<f32>):  // no predecessors
      %89 = mhlo.add %arg14, %arg15 : tensor<f32>
      "mhlo.return"(%89) : (tensor<f32>) -> ()
    }) {dimensions = dense<1> : tensor<1xi64>} : (tensor<128x10xf32>, tensor<f32>) -> tensor<128xf32>
    %40 = "mhlo.reshape"(%39) : (tensor<128xf32>) -> tensor<128x1xf32>
    %41 = mhlo.divide %40, %36 : tensor<128x1xf32>
    %42 = "mhlo.reduce"(%41, %10) ( {
    ^bb0(%arg14: tensor<f32>, %arg15: tensor<f32>):  // no predecessors
      %89 = mhlo.add %arg14, %arg15 : tensor<f32>
      "mhlo.return"(%89) : (tensor<f32>) -> ()
    }) {dimensions = dense<1> : tensor<1xi64>} : (tensor<128x1xf32>, tensor<f32>) -> tensor<128xf32>
    %43 = "mhlo.broadcast_in_dim"(%42) {broadcast_dimensions = dense<0> : tensor<1xi64>} : (tensor<128xf32>) -> tensor<128x10xf32>
    %44 = mhlo.multiply %43, %34 : tensor<128x10xf32>
    %45 = mhlo.add %37, %44 : tensor<128x10xf32>
    %46 = "mhlo.reduce"(%45, %10) ( {
    ^bb0(%arg14: tensor<f32>, %arg15: tensor<f32>):  // no predecessors
      %89 = mhlo.add %arg14, %arg15 : tensor<f32>
      "mhlo.return"(%89) : (tensor<f32>) -> ()
    }) {dimensions = dense<0> : tensor<1xi64>} : (tensor<128x10xf32>, tensor<f32>) -> tensor<10xf32>
    %47 = "mhlo.reshape"(%46) : (tensor<10xf32>) -> tensor<1x10xf32>
    %48 = "mhlo.reduce"(%47, %10) ( {
    ^bb0(%arg14: tensor<f32>, %arg15: tensor<f32>):  // no predecessors
      %89 = mhlo.add %arg14, %arg15 : tensor<f32>
      "mhlo.return"(%89) : (tensor<f32>) -> ()
    }) {dimensions = dense<0> : tensor<1xi64>} : (tensor<1x10xf32>, tensor<f32>) -> tensor<10xf32>
    %49 = "mhlo.dot_general"(%45, %24) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = ["DEFAULT", "DEFAULT"]} : (tensor<128x10xf32>, tensor<128x1024xf32>) -> tensor<10x1024xf32>
    %50 = "mhlo.transpose"(%49) {permutation = dense<[1, 0]> : tensor<2xi64>} : (tensor<10x1024xf32>) -> tensor<1024x10xf32>
    %51 = "mhlo.dot_general"(%45, %arg10) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [1]>, precision_config = ["DEFAULT", "DEFAULT"]} : (tensor<128x10xf32>, tensor<1024x10xf32>) -> tensor<128x1024xf32>
    %52 = "mhlo.select"(%25, %51, %11) : (tensor<128x1024xi1>, tensor<128x1024xf32>, tensor<128x1024xf32>) -> tensor<128x1024xf32>
    %53 = "mhlo.reduce"(%52, %10) ( {
    ^bb0(%arg14: tensor<f32>, %arg15: tensor<f32>):  // no predecessors
      %89 = mhlo.add %arg14, %arg15 : tensor<f32>
      "mhlo.return"(%89) : (tensor<f32>) -> ()
    }) {dimensions = dense<0> : tensor<1xi64>} : (tensor<128x1024xf32>, tensor<f32>) -> tensor<1024xf32>
    %54 = "mhlo.reshape"(%53) : (tensor<1024xf32>) -> tensor<1x1024xf32>
    %55 = "mhlo.reduce"(%54, %10) ( {
    ^bb0(%arg14: tensor<f32>, %arg15: tensor<f32>):  // no predecessors
      %89 = mhlo.add %arg14, %arg15 : tensor<f32>
      "mhlo.return"(%89) : (tensor<f32>) -> ()
    }) {dimensions = dense<0> : tensor<1xi64>} : (tensor<1x1024xf32>, tensor<f32>) -> tensor<1024xf32>
    %56 = "mhlo.dot_general"(%52, %18) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = ["DEFAULT", "DEFAULT"]} : (tensor<128x1024xf32>, tensor<128x1024xf32>) -> tensor<1024x1024xf32>
    %57 = "mhlo.transpose"(%56) {permutation = dense<[1, 0]> : tensor<2xi64>} : (tensor<1024x1024xf32>) -> tensor<1024x1024xf32>
    %58 = "mhlo.dot_general"(%52, %arg6) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [1]>, precision_config = ["DEFAULT", "DEFAULT"]} : (tensor<128x1024xf32>, tensor<1024x1024xf32>) -> tensor<128x1024xf32>
    %59 = "mhlo.select"(%19, %58, %11) : (tensor<128x1024xi1>, tensor<128x1024xf32>, tensor<128x1024xf32>) -> tensor<128x1024xf32>
    %60 = "mhlo.reduce"(%59, %10) ( {
    ^bb0(%arg14: tensor<f32>, %arg15: tensor<f32>):  // no predecessors
      %89 = mhlo.add %arg14, %arg15 : tensor<f32>
      "mhlo.return"(%89) : (tensor<f32>) -> ()
    }) {dimensions = dense<0> : tensor<1xi64>} : (tensor<128x1024xf32>, tensor<f32>) -> tensor<1024xf32>
    %61 = "mhlo.reshape"(%60) : (tensor<1024xf32>) -> tensor<1x1024xf32>
    %62 = "mhlo.reduce"(%61, %10) ( {
    ^bb0(%arg14: tensor<f32>, %arg15: tensor<f32>):  // no predecessors
      %89 = mhlo.add %arg14, %arg15 : tensor<f32>
      "mhlo.return"(%89) : (tensor<f32>) -> ()
    }) {dimensions = dense<0> : tensor<1xi64>} : (tensor<1x1024xf32>, tensor<f32>) -> tensor<1024xf32>
    %63 = "mhlo.dot_general"(%59, %arg0) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [0], rhs_contracting_dimensions = [0]>, precision_config = ["DEFAULT", "DEFAULT"]} : (tensor<128x1024xf32>, tensor<128x784xf32>) -> tensor<1024x784xf32>
    %64 = "mhlo.transpose"(%63) {permutation = dense<[1, 0]> : tensor<2xi64>} : (tensor<1024x784xf32>) -> tensor<784x1024xf32>
    %65 = mhlo.multiply %arg3, %9 : tensor<784x1024xf32>
    %66 = mhlo.add %65, %64 : tensor<784x1024xf32>
    %67 = mhlo.multiply %66, %8 : tensor<784x1024xf32>
    %68 = mhlo.subtract %arg2, %67 : tensor<784x1024xf32>
    %69 = mhlo.multiply %arg5, %5 : tensor<1024xf32>
    %70 = mhlo.add %69, %62 : tensor<1024xf32>
    %71 = mhlo.multiply %70, %4 : tensor<1024xf32>
    %72 = mhlo.subtract %arg4, %71 : tensor<1024xf32>
    %73 = mhlo.multiply %arg7, %7 : tensor<1024x1024xf32>
    %74 = mhlo.add %73, %57 : tensor<1024x1024xf32>
    %75 = mhlo.multiply %74, %6 : tensor<1024x1024xf32>
    %76 = mhlo.subtract %arg6, %75 : tensor<1024x1024xf32>
    %77 = mhlo.multiply %arg9, %5 : tensor<1024xf32>
    %78 = mhlo.add %77, %55 : tensor<1024xf32>
    %79 = mhlo.multiply %78, %4 : tensor<1024xf32>
    %80 = mhlo.subtract %arg8, %79 : tensor<1024xf32>
    %81 = mhlo.multiply %arg11, %3 : tensor<1024x10xf32>
    %82 = mhlo.add %81, %50 : tensor<1024x10xf32>
    %83 = mhlo.multiply %82, %2 : tensor<1024x10xf32>
    %84 = mhlo.subtract %arg10, %83 : tensor<1024x10xf32>
    %85 = mhlo.multiply %arg13, %1 : tensor<10xf32>
    %86 = mhlo.add %85, %48 : tensor<10xf32>
    %87 = mhlo.multiply %86, %0 : tensor<10xf32>
    %88 = mhlo.subtract %arg12, %87 : tensor<10xf32>
    return %68, %66, %72, %70, %76, %74, %80, %78, %84, %82, %88, %86 : tensor<784x1024xf32>, tensor<784x1024xf32>, tensor<1024xf32>, tensor<1024xf32>, tensor<1024x1024xf32>, tensor<1024x1024xf32>, tensor<1024xf32>, tensor<1024xf32>, tensor<1024x10xf32>, tensor<1024x10xf32>, tensor<10xf32>, tensor<10xf32>
  }
  func @predict(%arg0: tensor<128x784xf32>) -> tensor<128x10xf32> {
    %0 = iree_input.global.load @init_params$0 : tensor<784x1024xf32>
    %1 = iree_input.global.load @init_params$1 : tensor<1024xf32>
    %2 = iree_input.global.load @init_params$2 : tensor<1024x1024xf32>
    %3 = iree_input.global.load @init_params$3 : tensor<1024xf32>
    %4 = iree_input.global.load @init_params$4 : tensor<1024x10xf32>
    %5 = iree_input.global.load @init_params$5 : tensor<10xf32>
    %6 = call @main_12(%0, %1, %2, %3, %4, %5, %arg0) : (tensor<784x1024xf32>, tensor<1024xf32>, tensor<1024x1024xf32>, tensor<1024xf32>, tensor<1024x10xf32>, tensor<10xf32>, tensor<128x784xf32>) -> tensor<128x10xf32>
    return %6 : tensor<128x10xf32>
  }
  func private @main_12(%arg0: tensor<784x1024xf32>, %arg1: tensor<1024xf32>, %arg2: tensor<1024x1024xf32>, %arg3: tensor<1024xf32>, %arg4: tensor<1024x10xf32>, %arg5: tensor<10xf32>, %arg6: tensor<128x784xf32>) -> tensor<128x10xf32> {
    %0 = mhlo.constant dense<0.000000e+00> : tensor<f32>
    %1 = mhlo.constant dense<0xFF800000> : tensor<f32>
    %2 = mhlo.constant dense<0.000000e+00> : tensor<128x1024xf32>
    %3 = "mhlo.dot_general"(%arg6, %arg0) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = ["DEFAULT", "DEFAULT"]} : (tensor<128x784xf32>, tensor<784x1024xf32>) -> tensor<128x1024xf32>
    %4 = "mhlo.broadcast_in_dim"(%arg1) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<1024xf32>) -> tensor<1x1024xf32>
    %5 = "mhlo.broadcast_in_dim"(%4) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor<1x1024xf32>) -> tensor<128x1024xf32>
    %6 = mhlo.add %3, %5 : tensor<128x1024xf32>
    %7 = mhlo.maximum %6, %2 : tensor<128x1024xf32>
    %8 = "mhlo.dot_general"(%7, %arg2) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = ["DEFAULT", "DEFAULT"]} : (tensor<128x1024xf32>, tensor<1024x1024xf32>) -> tensor<128x1024xf32>
    %9 = "mhlo.broadcast_in_dim"(%arg3) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<1024xf32>) -> tensor<1x1024xf32>
    %10 = "mhlo.broadcast_in_dim"(%9) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor<1x1024xf32>) -> tensor<128x1024xf32>
    %11 = mhlo.add %8, %10 : tensor<128x1024xf32>
    %12 = mhlo.maximum %11, %2 : tensor<128x1024xf32>
    %13 = "mhlo.dot_general"(%12, %arg4) {dot_dimension_numbers = #mhlo.dot<lhs_contracting_dimensions = [1], rhs_contracting_dimensions = [0]>, precision_config = ["DEFAULT", "DEFAULT"]} : (tensor<128x1024xf32>, tensor<1024x10xf32>) -> tensor<128x10xf32>
    %14 = "mhlo.broadcast_in_dim"(%arg5) {broadcast_dimensions = dense<1> : tensor<1xi64>} : (tensor<10xf32>) -> tensor<1x10xf32>
    %15 = "mhlo.broadcast_in_dim"(%14) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor<1x10xf32>) -> tensor<128x10xf32>
    %16 = mhlo.add %13, %15 : tensor<128x10xf32>
    %17 = "mhlo.reduce"(%16, %1) ( {
    ^bb0(%arg7: tensor<f32>, %arg8: tensor<f32>):  // no predecessors
      %27 = mhlo.maximum %arg7, %arg8 : tensor<f32>
      "mhlo.return"(%27) : (tensor<f32>) -> ()
    }) {dimensions = dense<1> : tensor<1xi64>} : (tensor<128x10xf32>, tensor<f32>) -> tensor<128xf32>
    %18 = "mhlo.broadcast_in_dim"(%17) {broadcast_dimensions = dense<0> : tensor<1xi64>} : (tensor<128xf32>) -> tensor<128x1xf32>
    %19 = "mhlo.broadcast_in_dim"(%18) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor<128x1xf32>) -> tensor<128x10xf32>
    %20 = mhlo.subtract %16, %19 : tensor<128x10xf32>
    %21 = "mhlo.exponential"(%20) : (tensor<128x10xf32>) -> tensor<128x10xf32>
    %22 = "mhlo.reduce"(%21, %0) ( {
    ^bb0(%arg7: tensor<f32>, %arg8: tensor<f32>):  // no predecessors
      %27 = mhlo.add %arg7, %arg8 : tensor<f32>
      "mhlo.return"(%27) : (tensor<f32>) -> ()
    }) {dimensions = dense<1> : tensor<1xi64>} : (tensor<128x10xf32>, tensor<f32>) -> tensor<128xf32>
    %23 = "mhlo.broadcast_in_dim"(%22) {broadcast_dimensions = dense<0> : tensor<1xi64>} : (tensor<128xf32>) -> tensor<128x1xf32>
    %24 = "mhlo.log"(%23) : (tensor<128x1xf32>) -> tensor<128x1xf32>
    %25 = "mhlo.broadcast_in_dim"(%24) {broadcast_dimensions = dense<[0, 1]> : tensor<2xi64>} : (tensor<128x1xf32>) -> tensor<128x10xf32>
    %26 = mhlo.subtract %20, %25 : tensor<128x10xf32>
    return %26 : tensor<128x10xf32>
  }
}