Skip to content

Instantly share code, notes, and snippets.

@stellaraccident
Created November 26, 2021 23:43
Show Gist options
  • Save stellaraccident/c2f1ae7841c2a7fbc07c778874c50655 to your computer and use it in GitHub Desktop.
Save stellaraccident/c2f1ae7841c2a7fbc07c778874c50655 to your computer and use it in GitHub Desktop.
High Level JAX export
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>
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment