%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>

}

void BatchFunctionOp::eraseArguments(const BitVector& erase_indices) {
  const StringRef operand_segment_size_attr = getOperandSegmentSizeAttr();
  auto operandSegmentSizes = getOperation()->getAttrOfType<DenseI32ArrayAttr>(
      operand_segment_size_attr);

  // `operandSegmentSizes` attribute indicates the sizes of the two
  // variadic operands of `BatchFunctionOp`: `in_tensors` and

  func.return
}

// -----

func.func @testSendTPUEmbeddingGradients(%x: tensor<512x256xf32>) {
  "tf.SendTPUEmbeddingGradients"(%x) {N = 1 : i64, NN = 0 : i64, config = "", operandSegmentSizes = array<i32: 1, 0>} : (tensor<512x256xf32>) -> ()
  func.return
}

// -----

//===--------------------------------------------------------------------===//
//  tf.BatchToSpace

func.func @testXlaReduceToXlaVariadicReduceV2(%arg0: tensor<*xbf16>, %arg1: tensor<*xbf16>) -> tensor<*xbf16> {
  // CHECK: "tf.XlaVariadicReduceV2"(%arg0, %arg1) <{dimensions_to_reduce = [], operandSegmentSizes = array<i32: 1, 1>, reducer = @sum1}> {device = "/job:localhost/replica:0/task:0/device:GPU:0"} : (tensor<*xbf16>, tensor<*xbf16>) -> tensor<*xbf16>

    %2:4 = "tf_device.launch"() ({
      %4 = "tf.Cast"(%arg0) {Truncate = false} : (tensor<512xi64>) -> tensor<512xi32>
      %5 = "tf.TPUCopyWithDynamicShape"(%4,  %cst) {operandSegmentSizes = array<i32: 1, 1>} : (tensor<512xi32>, tensor<i32>) -> tensor<512xi32>
      tf_device.return %5 : tensor<512xi32>