func.return %1 : tensor<*xf32>

  // CHECK: "tf.BatchFunction"(%arg0, %0)
  // CHECK: operandSegmentSizes = array<i32: 1, 1>

            %2 = "tf.Cast"(%arg1) {Truncate = false} : (tensor<2048xi64>) -> tensor<2048xi32>
            %3:2 = "tf.TPUCopyWithDynamicShape"(%1, %2, %cst, %cst) {operandSegmentSizes = array<i32: 2, 2>} : (tensor<2048xi32>, tensor<2048xi32>, tensor<i32>, tensor<i32>) -> (tensor<2048xi32>, tensor<2048xi32>)
            // CHECK-NOT: tf.TPUAnnotateTensorsWithDynamicShape

      %outputs_6, %control_7 = tf_executor.island wraps "tf.ParseExample"(%outputs_4, %outputs, %outputs_0, %outputs_2) {dense_shapes = [#tf_type.shape<1>], device = "", operandSegmentSizes = array<i32: 1, 1, 0, 1, 1>, resultSegmentSizes = array<i32: 0, 0, 0, 1>} : (tensor<1x!tf_type.string>, tensor<1x!tf_type.string>, tensor<!tf_type.string>, tensor<i64>) -> tensor<*xi64>

func.func @serving_default(%arg0: tensor<1xf32>, %arg1: tensor<1xf32>) -> tensor<1xf32> {
  %0 = "tf.BatchFunction"(%arg0, %arg1) {f = @batched_func, num_batch_threads = 1 : i64, max_batch_size = 2 : i64, batch_timeout_micros = 10000 : i64, operandSegmentSizes = array<i32: 1, 1>} : (tensor<1xf32>, tensor<1xf32>) -> (tensor<1xf32>)
  return %0 : tensor<1xf32>
}
// The contents of `@serving_default` should have been inlined to `@batch_func`.

// CHECK-LABEL: func @replicate_derived_operandSegmentSizes
func.func @replicate_derived_operandSegmentSizes() {
  tf_device.replicate {n = 2 : i32, operandSegmentSizes = array<i32: 0, 0>} {
  }
  func.return

// CHECK:      tf_device.replicate
// CHECK-SAME: n = 2
// CHECK-NOT:  operandSegmentSizes
// CHECK-NEXT:   tf_device.return
}

// -----

// CHECK-LABEL: func @replicate_with_return

  %3 = "tf.Const"() {value = dense<0.0> : tensor<4x4xf32>} : () -> tensor<4x4xf32>
  "tf.SendTPUEmbeddingGradients"(%2, %3) {_tpu_embedding_layer = "call1", config = "\0A\0B\0C\0D", operandSegmentSizes = array<i32: 2, 0>} : (tensor<2x2xf32>, tensor<4x4xf32>) -> ()

  // CHECK: %[[CONST_MODE:.*]] = "tf.Const"() <{value = dense<"train"> : tensor<!tf_type.string>}> {_xla_outside_compilation = "0"} : () -> tensor<!tf_type.string>

  ) -> (tensor<1x1x!tf_type.string>, tensor<1x1x!tf_type.string>)
  func.return %4#0, %4#1 : tensor<1x1x!tf_type.string>, tensor<1x1x!tf_type.string>
// CHECK: "tf.ParseExample"(
// CHECK-SAME: operandSegmentSizes = array<i32: 1, 1, 0, 2, 2>, resultSegmentSizes = array<i32: 0, 0, 0, 2>
}

// CHECK-LABEL: TfMapDataset
func.func @TfMapDataset(%arg0: tensor<!tf_type.variant>) -> (tensor<!tf_type.variant>) {

        %0 = "tf.XlaLaunch"(%input) {function = @func, operandSegmentSizes = array<i32: 0, 1, 0>, device = "CPU:0"} : (tensor<f32>) -> tensor<f32>
        // expected-remark@above {{ID: 0}}
        // expected-remark@above {{Successors: {1}}}
        %1 = "tf.XlaLaunch"(%input) {function = @func, operandSegmentSizes = array<i32: 0, 1, 0>, device = "CPU:0"} : (tensor<f32>) -> tensor<f32>

    enable_large_batch_splitting = false, f = @batched_function,
    max_batch_size = 6 : i64, max_enqueued_batches = 10 : i64,
    num_batch_threads = 1 : i64, operandSegmentSizes = array<i32: 1, 0>,
    shared_name = "batch_function"
  } : (tensor<1xi32>) -> tensor<1xi32>

  // CHECK: [[z:%.*]] = tf_mlrt.executeop