registry.insert<quant::QuantizationDialect,
                    quantfork::QuantizationForkDialect>();
  }

  StringRef getArgument() const final {
    // This is the argument used to refer to the pass in
    // the textual format (on the commandline for example).
    return "tf-to-quant";
  }
  StringRef getDescription() const final {
    // This is a brief description of the pass.
    return "Legalize TF to quant ops dialect";

//   * has an attribute with the given `name`,
//   * and the attribute is an integer list of the form [1, X, Y, 1],
// and writes X, Y as 32-bit integer attribute to `x`, `y`.
bool TFIntListIs1XY1(Operation *op, StringRef name, IntegerAttr *x,
                     IntegerAttr *y) {
  auto attr = op->getAttrOfType<ArrayAttr>(name);
  if (!attr) return false;

  auto elements = attr.getValue();
  if (elements.size() != 4 ||

#ifndef TENSORFLOW_COMPILER_MLIR_TENSORFLOW_UTILS_EVAL_UTIL_H_
#define TENSORFLOW_COMPILER_MLIR_TENSORFLOW_UTILS_EVAL_UTIL_H_

#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/IR/Operation.h"  // from @llvm-project
#include "tensorflow/c/eager/c_api.h"

namespace tensorflow {

// Attempts to evaluates an MLIR Operation in TensorFlow eager mode with the

#ifndef TENSORFLOW_COMPILER_MLIR_TENSORFLOW_TRANSFORMS_BRIDGE_H_
#define TENSORFLOW_COMPILER_MLIR_TENSORFLOW_TRANSFORMS_BRIDGE_H_

#include <string>

#include "absl/base/attributes.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 "tensorflow/core/lib/core/status.h"

namespace mlir {
namespace TF {

namespace tfrt {
namespace fallback {

// Dialect for fallback operations.
class FallbackDialect : public Dialect {
 public:
  explicit FallbackDialect(MLIRContext *context);
  static StringRef getDialectNamespace() { return "tfrt_fallback"; }

  Type parseType(DialectAsmParser &parser) const override;
  void printType(Type type, DialectAsmPrinter &os) const override;
};

  ASSERT_TRUE(module);

  std::vector<std::pair<std::string, int>> func_and_index;
  ForEachArgConsumedByFallback(
      module.get(),
      [&func_and_index](llvm::StringRef func_name, int arg_index) {
        func_and_index.push_back({func_name.str(), arg_index});
      });

  ASSERT_EQ(func_and_index.size(), 1);
  EXPECT_EQ(func_and_index[0].first, "test");

==============================================================================*/

#include <memory>
#include <vector>

#include <gtest/gtest.h>
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Pass/Pass.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tf2xla/api/v2/legalize_tf.h"
#include "xla/client/client_library.h"
#include "xla/shape.h"
#include "xla/stream_executor/platform.h"

#include "llvm/ADT/APInt.h"
#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/Sequence.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringExtras.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/StringSwitch.h"
#include "llvm/ADT/iterator_range.h"
#include "llvm/Support/Casting.h"
#include "llvm/Support/FormatVariadic.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project

limitations under the License.
==============================================================================*/

#include "tensorflow/compiler/mlir/tensorflow/transforms/bridge.h"

#include "absl/log/log.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/Pass/PassManager.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"

// each time the pass runs. A special environment variable is set to locate the
// decomposition library.
class GraphDecomposePass : public MlirOptimizationPass {
 public:
  llvm::StringRef name() const override { return "tfr"; }

  // Whether to run this pass. If this is enabled, the GraphDef will be imported
  // to MLIR even no tf composition file is found.
  ::tensorflow::MlirOptimizationPassState GetPassState(