backprop.replaceAllUsesWith(slice_op.getResult());
}

// Checks if the input producer op is supported in this transform. Right now, we
// only check if it is a host tf.IteratorGetNext.
bool IsSupportedHostInputOp(Operation* op) {
  TF::IteratorGetNextOp iter = llvm::dyn_cast<TF::IteratorGetNextOp>(op);
  if (!iter) return false;
  auto device = op->getAttrOfType<StringAttr>(kDeviceAttr);
  if (!device) return false;