// extra attribute.
      if (!ops_with_tokens.contains(&op)) continue;

      if (!IsCommunicationOp(&op) && !SupportsCommunicationComputation(&op)) {
        return op.emitOpError(
            "does not support subcomputations with tf/xla communication ops");
      }

      op.setAttr(tensorflow::kXlaTokenInputNodesAttrName,
                 ArrayAttr::get(ctx, {token}));

// call to the reduce function in the ReduceDatasetOp.
LogicalResult DecomposeReduceDatasetInFunction(FuncOp function) {
  if (!llvm::hasSingleElement(function))
    return function.emitOpError("Expecting a single block function");

  auto decompose_result = function.walk([&](ReduceDatasetOp reduce_dataset) {
    if (!reduce_dataset->hasAttrOfType<StringAttr>(TF::kCompileDeviceTypeAttr))

      } else if (auto other_island_op =
                     owner->getParentOfType<tf_executor::IslandOp>()) {
        (*new_control_inputs)[other_island_op].push_back(sink_island_control);
      } else {
        owner->emitOpError("adding control dependency not supported");
        return signalPassFailure();
      }
    }
  }
  for (auto item :
       llvm::zip(island_op.getOutputs(), island_op.GetYield().getFetches()))

    if (!IsControlFlowOp(parent))
      return op->emitOpError()
             << "expects ancestor(s) to be of ['" << IfOp::getOperationName()
             << "', '" << func::FuncOp::getOperationName() << "']";

    if (!llvm::hasSingleElement(block->getParent()->getBlocks()))
      return op->emitOpError() << "expects single block region ancestor(s)";

    control_flow_ops.insert(parent);

        return;
      }
      iter_op = iter_op->getParentOp();
    }

    if (!IsSupportedOp(*op, supported_ops, tf_dialect)) {
      op->emitOpError() << "isn't compilable for TPU device. enable "
                           "soft_device_placement option to run on CPU";
      ++uncompilable_op_count;
    }
  });
  return uncompilable_op_count > 0;
}

    use_spmd = use_spmd_attr.getValue();

  auto num_cores_per_replica_attr =
      cluster_func->getAttrOfType<IntegerAttr>(kNumCoresPerReplicaAttr);
  if (!num_cores_per_replica_attr)
    return cluster_func.emitOpError(
        CreateMissingAttributeMsg(kNumCoresPerReplicaAttr));

  int num_cores_per_replica = num_cores_per_replica_attr.getInt();
  llvm::SmallVector<std::string> logical_device_vec(num_cores_per_replica);

    llvm::raw_string_ostream os(err_msg);
    expected_params.print(os);
    os << " vs. ";
    compared_params.print(os);
    os.flush();
    return op->emitOpError(err_msg);
  }
  return success();
}

quant::UniformQuantizedType GetFixedOutputRange(
    const bool is_signed, const int bit_width, const Type tensor_type,