#include <gmock/gmock.h>
#include <gtest/gtest.h>
#include "absl/status/status.h"
#include "absl/status/statusor.h"
#include "absl/strings/string_view.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/Support/raw_ostream.h"
#include "mlir/IR/DialectRegistry.h"  // from @llvm-project
#include "mlir/Pass/PassManager.h"  // from @llvm-project
#include "tensorflow/compiler/jit/xla_compile_util.h"

void FindCustomAggregatorOps(
    Region& region,
    const std::unordered_set<std::string>& aggregator_ops_to_ignore,
    SmallVector<Value>& statistics_outputs, SmallVector<StringRef>& ids,
    SmallVector<int32_t>& calibration_methods) {
  for (auto op : region.getOps<TF::CustomAggregatorOp>()) {
    if (aggregator_ops_to_ignore.count(op.getId().str())) continue;

    ids.push_back(op.getId());

#include "tensorflow/compiler/mlir/tf2xla/api/v2/tf_dialect_to_executor.h"

#include <memory>
#include <string>

#include "absl/log/log.h"
#include "absl/status/status.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/Pass/PassManager.h"  // from @llvm-project

#include "tensorflow/compiler/mlir/tf2xla/internal/legalize_tf_mlir.h"

#include <memory>
#include <string>
#include <vector>

#include "absl/log/log.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/IR/DialectRegistry.h"  // from @llvm-project
#include "mlir/IR/MLIRContext.h"  // from @llvm-project

  explicit InsertRestoreOpPass() = default;

  // The argument used to refer to the pass in the textual format (e.g. on the
  // commandline).
  StringRef getArgument() const final { return "quant-insert-restore-op"; }

  StringRef getDescription() const final {
    return "Creates RestoreV2 op to initialize the variables in the "
           "initializer function (`tf_saved_model.initializer_type == "

  // Dialect is used to model TF GradientDef. GetGradientAttrName() returns the
  // string description of gradient attribute.
  static StringRef GetGradientAttrName() { return "tf.gradient"; }

  // This attribute marks if a function is stateful.
  // Returns the string description of stateful attribute.
  static StringRef GetStatefulAttrName() { return "tf.signature.is_stateful"; }

  // Returns true if the op can be duplicated during transformations.

  }
}

LogicalResult DeviceTarget::RegisterKernel(
    llvm::StringRef kernel, const KernelSpecs::Signature& signature,
    const ScaleFn& fn, const ScaleDecomposeFn& dfn) {
  return specs_[kernel].Add(signature, {ScaleConstraintType::CustomScale, fn});
}

namespace ph = std::placeholders;

LogicalResult DeviceTarget::RegisterKernel(
    llvm::StringRef kernel, const KernelSpecs::Signature& signature,

#define TENSORFLOW_COMPILER_MLIR_TF2XLA_TRANSFORMS_LEGALIZE_TF_WITH_TF2XLA_PASSES_H_

#include <memory>
#include <optional>

#include "llvm/ADT/StringRef.h"
#include "mlir/IR/MLIRContext.h"  // from @llvm-project
#include "mlir/IR/PatternMatch.h"  // from @llvm-project
#include "mlir/Pass/Pass.h"  // from @llvm-project
#include "mlir/Pass/PassRegistry.h"  // from @llvm-project

// variables. If the `saved_model_dir` directory path is provided, then the
// `tf_saved_model.asset` ops will be freezed.
void AddTFToTFLConversionPasses(llvm::StringRef saved_model_dir,
                                const toco::TocoFlags& toco_flags,
                                const mlir::TFL::PassConfig& pass_config,
                                mlir::OpPassManager* pass_manager);

  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(CastBf16OpsToF32Pass)
  explicit CastBf16OpsToF32Pass() = default;

  StringRef getArgument() const final {
    // This is the argument used to refer to the pass in
    // the textual format (on the commandline for example).
    return "quant-cast-bf16-ops-to-f32";
  }
  StringRef getDescription() const final {
    return "Cast BF16 operations to F32.";
  }

  void runOnOperation() override;