func.return %1 : tensor<*xf32>
}

// -----

// The num_batch_threads remains 3 (the same as the original attribute)

               << ") to be equal to 'n' (" << n << ")";
    }
  }

  Block& block = op.getBody().front();

  auto operandSegmentSizes = op.getProperties().operandSegmentSizes;
  const int32_t num_replicated_inputs = operandSegmentSizes[0];
  const int32_t num_packed_inputs = operandSegmentSizes[1];

  if (num_replicated_inputs % n != 0)
    return op.emitOpError()

  func.return %1 : tensor<*xf32>

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

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

    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

    data_format = "NHWC", dilations = [1, 1, 1, 1], epsilon = 9.99999974E-5 : f32, explicit_paddings = [], filter_format = "HWIO", fused_ops = ["BiasAdd", "Relu"], leakyrelu_alpha = 2.000000e-01 : f32, num_args = 2 : i64, operandSegmentSizes = array<i32: 1, 1, 2, 2>, padding = "SAME", strides = [1, 1, 1, 1], use_cudnn_on_gpu = true
    } : (tensor<1x300x300x40xi8>, tensor<3x3x40x40xi8>, tensor<40xf32>, tensor<0xi8>, tensor<f32>, tensor<f32>) -> tensor<1x300x300x40xi8>

    %1 = func.call @finalize(%arg4, %arg5) : (tensor<6xf32>, tensor<2xf32>) -> tensor<5xf32>
    "tf.Yield"(%1) : (tensor<5xf32>) -> ()
  }) {device = "/job:tpu_host_worker/replica:0/task:0/device:CPU:0", metadata = "", operandSegmentSizes = array<i32: 1, 2, 1>, output_shapes = [#tf_type.shape<>], output_types = [!tf_type.string]} : (tensor<4xf32>, tensor<3xf32>, tensor<!tf_type.resource>, tensor<2xf32>) -> tensor<!tf_type.variant>
  return
}

// -----

    "tf.SendTPUEmbeddingGradients"(%2#0, %2#1) {_tpu_embedding_layer = "call1", config = "\0A\0B\0C\0D", operandSegmentSizes = array<i32: 2, 0>} : (tensor<2x2xf32>, tensor<4x4xf32>) -> ()
    tf_device.return
  }) {allow_soft_placement = true, num_cores_per_replica = 1, topology =  "", device_assignment =  []} : () -> ()
  func.return
}

}

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