#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/lite/stablehlo/transforms/passes.h.inc"

class TransposeCommuteOpsPass
    : public impl::TransposeCommuteOpsPassBase<TransposeCommuteOpsPass> {
  void runOnOperation() override;
};

// Inversely permutate a given vector
static SmallVector<int64_t> InvertPermutationToVector(ArrayRef<int64_t> vec,
                                                      ArrayRef<int64_t> perm) {

 public:
  using impl::FoldConstantTransposePassBase<
      FoldConstantTransposePass>::FoldConstantTransposePassBase;

 private:
  void runOnOperation() override;
};

void FoldConstantTransposePass::runOnOperation() {
  func::FuncOp func_op = getOperation();
  MLIRContext& ctx = getContext();

  RewritePatternSet patterns(&ctx);
  patterns.add<FoldTransposedConstantOp>(&ctx);

           "module for post-training dynamic range case";
  }

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

  void runOnOperation() override;

 private:
  Option<OpSet> target_opset_{
      *this, "target-opset", llvm::cl::init(OpSet::TF),
      llvm::cl::desc("Choose target opset."),
      llvm::cl::values(

#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) {

    : public impl::TensorFlowShapeInferencePassBase<ShapeInference> {
 public:
  ShapeInference() = default;
  explicit ShapeInference(ArrayRef<ArrayRef<int64_t>> input_shapes)
      : input_shapes_(input_shapes) {}
  void runOnOperation() override {
    // Parse `input_arg_shapes_` if provided (test only)
    SmallVector<ArrayRef<int64_t>> input_shapes_vec;
    absl::StatusOr<SmallVector<SmallVector<int64_t>>> parsed_shapes;

           "`tf.AssignVariableOp(tf.VarHandleOp, restore_op_output#N)`, where "
           "`restore_op_output#N` is the Nth output of the newly created "
           "RestoreV2Op.";
  }

  void runOnOperation() override;
};

// Finds `tf.AssignVariableOp(tf.VarHandleOp, tf.Const)` patterns and returns
// the `tf.VarHandleOp`s that are initialized by these `tf.AssignVariableOp`s.
std::vector<TF::VarHandleOp> CollectVariableOps(

  explicit FreezeGlobalTensorsPass(bool allow_mutable_tensors) {
    this->allow_mutable_tensors = allow_mutable_tensors;
  }
  void runOnOperation() override;
};

void FreezeGlobalTensorsPass::runOnOperation() {
  auto module = getOperation();
  if (!tf_saved_model::HasTfSavedModelSemantics(module)) return;

  DataFlowSolver solver;
  solver.load<dataflow::DeadCodeAnalysis>();

  explicit LiftQuantizableSpotsAsFunctionsPass(
      QuantizationSpecs quantization_specs)
      : quantization_specs_(std::move(quantization_specs)) {}

 private:
  void runOnOperation() override;

  // No explicit quantization spec is specified by default. Implicitly this
  // means that all quantizable units will be identified and lifted.
  QuantizationSpecs quantization_specs_{};
};