auto sharding_attr = mlir::dyn_cast<StringAttr>(attr);
  if (!sharding_attr)
    return op->emitOpError(
        llvm::formatv(kBadStringArrayElementMsg, name, index));
  if (tensorflow::DecodeShardingAttribute(sharding_attr, *sharding_ptr)
          .failed()) {
    return op->emitOpError(llvm::formatv(kBadArrayElementMsg, name, index,
                                         sharding_attr.getValue(),

  if (!op->hasAttr(TF::kTpuReplicateAttr)) {
    op->emitOpError("TF2XLA TPU bridge input check: " + errormsg() +
                    "missing _tpu_replicate attr");
    return false;
  }
  auto opattr = op->getAttr(TF::kTpuReplicateAttr);
  if (opattr != attr) {
    op->emitOpError("TF2XLA TPU bridge input check: " + errormsg() +
                    "invalid _tpu_replicate attr.")

        return fetch.emitOpError()
               << "operand #" << i
               << " is a control type, can't be bound to a graph result";
      break;
    }
    if (i >= num_results)
      return fetch.emitOpError()
             << "operand #" << i << " does not have a graph results to bind";
    if (graph.getResult(i).getType() != operand.getType()) {
      return fetch.emitOpError()

  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 "

  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(

        if (!backend_config) {
          op->emitOpError()
              << "is missing attribute '" << kTfBackendConfigAttrName << "'";
          return WalkResult::interrupt();
        }

        auto called_index_attr = mlir::dyn_cast_or_null<IntegerAttr>(
            backend_config.get(kCalledIndexAttrName));
        if (!called_index_attr) {
          op->emitOpError()