// CHECK:      "tf.EnqueueTPUEmbeddingArbitraryTensorBatch"
// CHECK-SAME:   _xla_replica_id = 1
// CHECK:      "tf.A"
// CHECK-NOT:   _xla_replica_id
// CHECK:      tf_executor.fetch


// Tests tf._TPUDeviceOrdinalPlaceholder ops are replaced with explicit device
// ordinal constant values based on the first TPU core device id.
// CHECK-LABEL: func @device_ordinals
func.func @device_ordinals() {
  tf_executor.graph {

    return region && region->isProperAncestor(replicate_region);
  };

  for (Value operand : op->getOperands())
    if (!ancestor_of_replicate(operand.getParentRegion())) return false;

  // _TPUDeviceOrdinalPlaceholder implicitly depends on the replica.
  if (llvm::isa<TF::_TPUDeviceOrdinalPlaceholderOp>(op)) return false;

  bool has_replicate_operands = false;