if (!replicate_attr || !compile_attr)
    return op.emitOpError() << "is expected to have either both or none of '"
                            << kReplicationInfoAttr << "' and '"
                            << kCompileDeviceTypeAttr << "' attributes.";
  if (replicate_attr.getValue().empty())
    return op.emitOpError()
           << "has an empty '" << kReplicationInfoAttr << "' attribute.";

  if (!matchPattern(group_assignment_value, m_Constant(&group_assignment))) {
    return op->emitOpError() << "expects constant group_assignment";
  }
  replica_groups = mlir::cast<DenseIntElementsAttr>(
      hlo::convertElementsAttr(group_assignment, builder.getIntegerType(64)));
  if (replica_groups.getType().getRank() != 2) {
    return op->emitOpError() << "group_assignment should have rank 2, got "

      // Invalid case since there are multiple loc:@ attributes.
      op->emitOpError()
          << "expects only one named location in '_class' attribute, but got "
          << classes_attr;
      return StringRef();
    }
    result = node_name.drop_front(kLocationPrefix.size());
  }
  if (result.empty()) {
    op->emitOpError() << "expects variable name in '_class' attribute, but got "

      /*target_tensor_names=*/{}, &resource_tensors);
  if (!status.ok()) {
    return module.emitOpError()
           << "failed to run the provided session: " << status.message();
  }

  const DeviceMgr* device_manager;
  if (!(session->LocalDeviceManager(&device_manager).ok())) {
    return module.emitOpError() << "failed to get local device manager";
  }

  // Read all underlying tensors of the variables from the session.

    if (!string_attr)
      return op->emitOpError(llvm::formatv(
          "bad '{0}' attribute at index {1}, not a string", kDevicesAttr, idx));

    if (DeviceNameUtils::ParseFullName(string_attr.getValue().str(), &device)) {
      devices->AddDevice(device);
    } else {
      return op->emitOpError(
          llvm::formatv("bad '{0}' attribute, '{1}', not a valid device",

  if (it == buffer_to_size->end()) {
    return scatter.emitOpError("unknown tensor list");
  }
  auto buffer = scatter.getInputHandle();
  OpBuilder builder(scatter);
  auto indices_type =
      mlir::cast<RankedTensorType>(scatter.getIndices().getType());
  if (!indices_type) return scatter.emitOpError("unranked indices shape");
  auto shape_type = RankedTensorType::get({2}, builder.getIntegerType(32));

    if (IfOp if_op = llvm::dyn_cast<IfOp>(op)) {
      if (failed(ConvertIfOp(if_op))) {
        op->emitOpError() << "failed to convert to region form";
        return WalkResult::interrupt();
      }
    } else if (CaseOp case_op = llvm::dyn_cast<CaseOp>(op)) {
      if (failed(ConvertCaseOp(case_op))) {
        op->emitOpError() << "failed to convert to region form";
        return WalkResult::interrupt();
      }

  if (!lengths_const) return split.emitOpError("non-constant split lengths");
  *count = lengths_const.getValue().getNumElements();
  if (*count <= 0) return split.emitOpError("non-positive split count");
  auto buffer_type = split.getValue().getType().dyn_cast<RankedTensorType>();
  if (!buffer_type || !buffer_type.hasStaticShape() ||
      buffer_type.getRank() < 1) {
    return split.emitOpError("unknown or invalid split tensor shape");
  }

  for (OpType op : func.getOps<OpType>()) {
    if (found_op != nullptr) {
      func.emitOpError() << "number of " << found_op.getOperationName()
                         << " in loop body is not 1";
      return LogicalResult::failure();
    }
    if (GetReplicationAttr(op) != replication_attr) {
      op.emitOpError() << "is not part of the replication region "

        old_partitioned_input.emitOpError(
            "TPUPartitionedInputV2 variable used in outside compiled code is "
            "only supported with REPLICATED sharding");
        return failure();
      } else if (is_packed) {
        // TODO(b/202047549) Generalize to the packed case.
        old_partitioned_input.emitOpError(