Operation* defining_op = value.getDefiningOp();
    return defining_op && ops_depend_on_cluster.contains(defining_op);
  });
}
}  // namespace

llvm::StringMap<SmallVector<Cluster>> BuildAllClusters(
    Block& block, const TF::SideEffectAnalysis::Info& side_effect_analysis,
    std::function<std::string(Operation*)> get_target,
    std::function<bool(Operation*)> is_ignored_op) {

  auto compile =
      llvm::cast<TF::_TPUCompileMlirOp>(compile_launch.GetBody().front());
  tensorflow::tpu::TPUCompileMetadataProto metadata;
  metadata.ParseFromString(compile.getMetadata().str());
  llvm::SmallVector<llvm::SmallVector<int64_t, 4>, 4> input_mappings =
      tensorflow::GetMetadataArgumentMapping(metadata);

  bool metadata_updated = false;
  auto maybe_replicate =

        log_dir_path_(std::move(log_dir_path)) {}

 private:
  SmallVector<NamedAttribute> CreateDumpAttributes(
      PatternRewriter &rewriter, const StringRef folder_name,
      const StringRef file_name, const bool enabled, const StringRef func_name,
      const StringRef node_name) const {
    SmallVector<NamedAttribute> dump_attributes{
        rewriter.getNamedAttr("log_dir_path",

    return op->emitError() << "cannot partition " << first_operand_type
                           << " (rank = " << rank << ") along dimension "
                           << partition_dim << ".";
  }

  llvm::SmallVector<int64_t, 4> partition_dims(is_replicated ? 0 : rank, 1);
  if (!is_replicated) {
    partition_dims[partition_dim] = num_cores_per_replica.getInt();
  }

  if constexpr (is_input) {

==============================================================================*/
#include "tensorflow/compiler/mlir/tensorflow/transforms/mark_initialized_variables.h"

#include <string>
#include <vector>

#include "llvm/ADT/SmallVector.h"
#include "mlir/IR/Block.h"  // from @llvm-project
#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
#include "mlir/IR/MLIRContext.h"  // from @llvm-project
#include "mlir/IR/Threading.h"  // from @llvm-project

// that can be used with InferShapeForFunction.
// TF NodeShape uses `,` to separate dimensions, and `:` to separate arguments.
// Ex: 1,2:3,4,5:6,? --> [[1, 2], [3, 4, 5], [6, ?]]
absl::StatusOr<SmallVector<SmallVector<int64_t>>> ParseArgumentShapes(
    absl::string_view input_shapes);

// Given a list of refined shapes matching the function arguments of func, runs

  if (!attr) return success();

  // ID pairs are separated by `,`.
  llvm::SmallVector<std::string, 8> str_list =
      absl::StrSplit(attr.getValue().str(), ',', absl::SkipWhitespace());
  id_pairs.reserve(str_list.size());
  for (const std::string& str : str_list) {
    // IDs of one pair are separated by `:`.
    llvm::SmallVector<std::string, 8> id_pair = absl::StrSplit(str, ':');

    // Check for malformed IDs.

    if (dyn_cast_or_null<TF::CustomAggregatorOp>(op)) return failure();

    // The CustomAggregatorOp is only added after quantizable values.
    SmallVector<Value> quantizable_values;
    SmallVector<std::string> aggregator_ids;
    if (IsCallToQuantizableLiftedFunction(op)) {
      std::optional<StringRef> composite_function_name =
          GetCompsiteFunctionName(op);

    } else if (auto qtype =
                   cast_qtype.dyn_cast<quant::UniformQuantizedPerAxisType>()) {
      SmallVector<float> scales(qtype.getScales().begin(),
                                qtype.getScales().end());
      SmallVector<int32_t> zps(qtype.getZeroPoints().begin(),
                               qtype.getZeroPoints().end());
      const size_t num_channels = qtype.getScales().size();

    ArrayRef<int64_t> input_shape = input_tensor.getShapedType().getShape();
    auto output_type = mlir::cast<ShapedType>(op.getOutput().getType());

    SmallVector<int32_t, 4> perm;
    SmallVector<int64_t, 4> output_shape;
    for (int i = 0; i < num_dimensions; ++i) {
      perm.push_back(perm_tensor.getValues<IntegerAttr>()[i].getInt());
      output_shape.push_back(input_shape[perm[i]]);