namespace {

using ::mlir::stablehlo::ConvDimensionNumbersAttr;

class NchwConvolutionToNhwcPass
    : public impl::NchwConvolutionToNhwcPassBase<NchwConvolutionToNhwcPass> {
 private:
  void runOnOperation() override;
};

// Rewrites NCHW convolution to NHWC.
// * Src dimension numbers: [b, f, 0, 1]x[o, i, 0, 1]->[b, f, 0, 1]
// * Dst dimension numbers: [b, 0, 1, f]x[0, 1, i, o]->[b, 0, 1, f]

    return "quant-add-quantization-unit-loc";
  }
  StringRef getDescription() const final {
    return "Add QuantizationUnitLoc to quantizable layers.";
  }

 private:
  void runOnOperation() override;
};

// TF graph nodes are imported with one of following location patterns:
//   FusedLoc[NameLoc(op_type:), ..., NameLoc(node_name@func_name)] or

#define GEN_PASS_DEF_CLUSTERFORMATIONPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

struct ClusterFormationPass
    : public impl::ClusterFormationPassBase<ClusterFormationPass> {
  void runOnOperation() override;
};

void ReplaceLiveOutExternalUses(llvm::ArrayRef<Value> live_outs,
                                tf_device::LaunchOp launch_op) {
  Region* launch_op_region = &launch_op.getBody();

#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

namespace {
struct MarkInputOutputAliasesPass
    : public impl::MarkInputOutputAliasesPassBase<MarkInputOutputAliasesPass> {
  void runOnOperation() override;
};

constexpr char kAliasingAttr[] = "tf.aliasing_output";
constexpr int kUnassigned = -1;

struct AliasInfo {
  AliasInfo() : input_index(kUnassigned), output_index(kUnassigned) {}

class TPUCleanupClusterAttributesPass
    : public impl::TPUCleanupClusterAttributesPassBase<
          TPUCleanupClusterAttributesPass> {
 public:
  void runOnOperation() override {
    auto traverse_op = [&](Operation* op, tf_device::ClusterOp tpu_cluster) {
      if (isa<tf_device::ClusterOp>(op)) return WalkResult::advance();
      op->removeAttr(TF::kReplicationInfoAttr);

#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

struct ReplicateInvariantOpHoistingPass
    : public impl::ReplicateInvariantOpHoistingPassBase<
          ReplicateInvariantOpHoistingPass> {
  void runOnOperation() override;
};

// Check if op directly uses a key in `virtual_devices`.
bool DirectUseOfVirtualDevice(const DictionaryAttr& virtual_devices,
                              Operation* op) {

      : accepted_dialects_(accepted_dialects), os_(os), total_ops_(0) {}

  // Prints the resultant operation statistics pos_t iterating over the module.
  void runOnOperation() override;

  // Prints summary of op stats.
  void PrintSummary();

  // Keeps track of dtype counts per op.
  void CountOp(DenseMap<StringRef, llvm::StringMap<int64_t>> &op_count_map,

  }

  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<TF::TensorFlowDialect, func::FuncDialect>();
  }

 private:
  void runOnOperation() override;

  // Returns the function library for the given quantization method and opset
  // pair.
  llvm::StringRef GetFunctionLibrary(QuantMethod quantization_method,

    registry.insert<TF::TensorFlowDialect>();
  }

 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MockSuccessPass)

  explicit MockSuccessPass() = default;

 private:
  void runOnOperation() override {
    getOperation().walk([](Operation* nestedOp) {
      nestedOp->emitError()
          << "Error at " << nestedOp->getName().getStringRef().str() << " op";
    });
  };
};

class ExtractTPUCopyWithDynamicShapeOpPass
    : public impl::ExtractTPUCopyWithDynamicShapeOpPassBase<
          ExtractTPUCopyWithDynamicShapeOpPass> {
  void runOnOperation() override;
};

// Finds op that created a given value. If the value is a BlockArgument, this
// returns the owner of the Block.
Operation* GetOpOfValue(Value value) {