// CHECK:           %[[VAL_2:.*]] = stablehlo.tuple %[[VAL_1]] : tuple<tensor<3x4xf32>>
// CHECK:           return %[[VAL_2]] : tuple<tensor<3x4xf32>>
// CHECK:         }
func.func @tfl_wrapped_jax_random_normal(%arg0: tensor<2xui32>) -> tuple<tensor<3x4xf32>> {
  // This is a fake jax random normal body.

# CHECK:     %[[VAL_2:.*]] = "tfl.batch_matmul"(%[[VAL_0]], %[[VAL_1]]) <{adj_x = false, adj_y = false}> : (tensor<2x5x3xf32>, tensor<3x7xf32>) -> tensor<2x5x7xf32>
# CHECK:     return %[[VAL_2]] : tensor<2x5x7xf32>

// CANON:           [[VAL_1:%.*]] = arith.constant dense<1.000000e+00> : tensor<256x256xf32>
// CANON:           [[VAL_2:%.*]] = arith.constant dense<0> : tensor<i32>
// CANON:           [[VAL_6:%.*]]:3 = "tfl.while"([[VAL_2]], [[VAL_2]], [[VAL_0]]) <{is_stateless = true}> ({
// CANON:           ^bb0([[VAL_7:%.*]]: tensor<*xi32>, [[VAL_8:%.*]]: tensor<*xi32>, [[VAL_9:%.*]]: tensor<*xf32>):

# CHECK-DAG:       %[[VAL_2:.*]] = arith.constant dense<[1, 0]> : tensor<2xi32>
# CHECK-DAG:       %[[VAL_3:.*]] = "tfl.no_value"() <{value}> : () -> none
# CHECK-DAG:       %[[VAL_6:.*]] = arith.constant dense<0> : tensor<i32>

// CHECK:  [[VAL_2:%.*]] = "tfl.concatenation"([[VAL_0]], [[VAL_1]]) <{axis = 3 : i32, fused_activation_function = "NONE"}> {tac.device = "GPU", tac.inference_type = "FLOAT"} : (tensor<1x1x1x1xf32>, tensor<1x1x1x1xf32>) -> tensor<1x1x1x2xf32>
// CHECK:  [[VAL_3:%.*]] = "tfl.reshape"([[VAL_2]], %{{.*}}) : (tensor<1x1x1x2xf32>, tensor<2xi32>) -> tensor<2x1xf32>