read_use.set(new_reads[0]);
    }
  }
  return success();
}

namespace {
constexpr char kDeviceAttr[] = "device";
constexpr char kFuncDeviceAttr[] = "tf.device";
mlir::Attribute GetDeviceOfResource(mlir::func::FuncOp func,
                                    mlir::Value resource) {
  if (auto* resource_op = resource.getDefiningOp()) {
    return resource_op->getAttr(kDeviceAttr);
  } else {