%1 = "tfl.pseudo_const"() { value = dense<1.0> : tensor<16xf32>} : () -> tensor<16xf32>
  %2 = "tfl.conv_3d"(%arg0, %0, %1) {dilation_h_factor = 1 : i32, dilation_w_factor = 1 : i32, dilation_d_factor = 1 : i32, fused_activation_function = "NONE", padding = "SAME", stride_h = 1 : i32, stride_w = 1 : i32, stride_d = 1 : i32} : (tensor<1x32x32x32x8xf32>, tensor<1x1x1x8x16x!quant.uniform<i8:f32, 1.0>>, tensor<16xf32>) -> tensor<1x32x32x32x16xf32>

// CHECK: %0 = "tfl.pseudo_const"() <{value = dense<2> : tensor<2xi32>}> : () -> tensor<2xi32>
// CHECK: %1 = "tfl.reshape"(%arg0, %0) : (tensor<1x4xf32>, tensor<2xi32>) -> tensor<2x2xf32>

{
  "version": 3,
  "operator_codes": [
    {
      "builtin_code": "RESHAPE"
    }
  ],
  "subgraphs": [
    {
      "tensors": [
        {
          "shape": [1, 4],

func.func @add_with_activation_transpose_rank_two(%arg0: tensor<1x2xf32>) -> tensor<2x1xf32> {
  %0 = stablehlo.constant dense<2.000000e+00> : tensor<2x1xf32>
  %1 = stablehlo.transpose %arg0, dims = [1, 0] : (tensor<1x2xf32>) -> tensor<2x1xf32>
  %2 = stablehlo.add %1, %0 : tensor<2x1xf32>
  return %2 : tensor<2x1xf32>
}
// CHECK: %[[TRANSPOSE_0:.+]] = stablehlo.transpose

}
func.func @_func(%arg0: tensor<2x4xf32>, %arg1: tensor<4x2xf32>) -> tensor<2x2xf32> {
  %0 = "tf.MatMul"(%arg0, %arg1) {_XlaSharding = "\08\03\1A\02\02\01\22\02\00\01"} : (tensor<2x4xf32>, tensor<4x2xf32>) -> tensor<2x2xf32>
  %1 = "tf.Identity"(%0) : (tensor<2x2xf32>) -> tensor<2x2xf32>
  return %1 : tensor<2x2xf32>
}

// -----
// The following op sharding is used in the following test case:

// CHECK-LABEL: func @testWhileCasts(%arg0: tensor<!tf_type.variant<tensor<1x3xf32>>>) -> tensor<!tf_type.variant<tensor<*xf32>>>
func.func @testWhileCasts(%arg0: tensor<!tf_type.variant<tensor<1x3xf32>>>) -> (tensor<!tf_type.variant<tensor<*xf32>>>) {
  %0 = "tf.While"(%arg0) {
    cond = @testWhileCond, body = @testWhileBody, is_stateless = false
  } : (tensor<!tf_type.variant<tensor<1x3xf32>>>) -> (tensor<!tf_type.variant<tensor<*xf32>>>)

        return %2 : tensor<?x2xf32>
      }
      func.func private @composite_fn_1(%arg0: tensor<?x2xf32>, %arg1: tensor<2x2xf32>, %arg2: tensor<2xf32>) -> tensor<?x2xf32> attributes {_from_xla_call_module, tf_quant.composite_function} {
        return %arg0 : tensor<?x2xf32>
      }
    }
  )mlir";

  %4 = "tf.MatMul"(%arg0, %3) {device = "", transpose_a = false, transpose_b = false} : (tensor<2x3xf32>, tensor<3x4xf32>) -> tensor<2x4xf32>
  %5 = "tf.Identity"(%4) {device = ""} : (tensor<2x4xf32>) -> tensor<2x4xf32>
  %6 = "tf.Identity"(%5) {device = ""} : (tensor<2x4xf32>) -> tensor<2x4xf32>
  func.return %6 : tensor<2x4xf32>

  %2:2 = "tf.IteratorGetNext"(%arg0) {device = "/device:CPU:0"} : (tensor<*x!tf_type.resource>) -> (tensor<128xf32>, tensor<128xf32>)
  %3:2 = "tf.IteratorGetNext"(%arg1) {device = "/device:CPU:0"} : (tensor<*x!tf_type.resource>) -> (tensor<128xf32>, tensor<128xf32>)
  // CHECK: "tf.TPUCompileSucceededAssert"(%[[COMPILE]]#0)
  "tf_device.launch"() ({

  func.return %8: tensor<4x4xf32>
}

// CHECK-LABEL: testSubOfNeg
func.func @testSubOfNeg(%arg0: tensor<8x16xf32>, %arg1: tensor<8x16xf32>) -> tensor<8x16xf32> {
  %0 = "tf.Neg"(%arg1) : (tensor<8x16xf32>) -> tensor<8x16xf32>
  %1 = "tf.Sub"(%arg0, %0) {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<8x16xf32>, tensor<8x16xf32>) -> tensor<8x16xf32>
  func.return %1: tensor<8x16xf32>

  %one = "glue.constant"() { value = 1: i32 } : () -> i32
  %done = "glue.compare" (%one, %one) { predicate = #glue<"compare LTE"> } : (i32, i32) -> i1
  %2 = mhlo.constant dense<[[1.1]]> : tensor<1x1xf32>
  %3 = mhlo.multiply %2, %2 : tensor<1x1xf32>
  %cst = "tf.Const"() {value = dense<0.0> : tensor<f32>} : () -> tensor<f32>
  %0 = "tf.AddV2"(%arg0, %cst) {device = "/device:CPU:0"} : (tensor<f32>, tensor<f32>) -> tensor<f32>