}

TensorFlowDialect::ConstantFoldHook TensorFlowDialect::constant_fold_hook_;

TensorFlowDialect::TensorFlowDialect(MLIRContext *context)
    : Dialect(/*name=*/"tf", context, TypeID::get<TensorFlowDialect>()) {
  context->getOrLoadDialect<::mlir::tf_type::TFTypeDialect>();
  addOperations<
#define GET_OP_LIST
#include "tensorflow/compiler/mlir/tensorflow/ir/tf_all_ops.cc.inc"
      >();

#include "tensorflow/compiler/mlir/tensorflow/ir/tf_types.h"

namespace mlir {
namespace TF {

class TensorFlowRegistryEffectInterfaceFallback;

class TensorFlowDialect final : public Dialect {
 public:
  explicit TensorFlowDialect(MLIRContext *context);
  ~TensorFlowDialect() override;

  static StringRef getDialectNamespace() { return "tf"; }

  // Overrides to redirect to tf_type dialect.

    that disperse values. (ex: convolution, dot_general)
  }];
  let dependentDialects = [
      "mlir::func::FuncDialect",
      "mlir::stablehlo::StablehloDialect",
      "TF::TensorFlowDialect",
  ];
}

def ReplaceStablehloOpsInMainFunctionWithXlaCallModuleOpsPass : Pass<"stablehlo-replace-stablehlo-ops-in-main-function-with-xla-call-module-ops", "mlir::ModuleOp"> {

  let description = [{
    Runs the pre calibration passes for post-training quantization with default
    configuration.
  }];
  let dependentDialects = [
    "mlir::stablehlo::StablehloDialect", "mlir::TF::TensorFlowDialect",
    "mlir::func::FuncDialect", "mlir::tf_executor::TensorFlowExecutorDialect",
    "mlir::mhlo::MhloDialect", "mlir::vhlo::VhloDialect",
  ];
}

}

func.func @func(%arg0: tensor<i32>) -> tensor<i32> {
  func.return %arg0 : tensor<i32>
}
)mlir";
  mlir::MLIRContext context;
  context.loadDialect<mlir::func::FuncDialect, mlir::TF::TensorFlowDialect>();
  mlir::OwningOpRef<mlir::ModuleOp> module =
      mlir::parseSourceString<mlir::ModuleOp>(code, &context);
  ASSERT_TRUE(module);
  auto entry_funcs = GetEntryFunctions(*module);
  EXPECT_EQ(entry_funcs.size(), 2);

  let description = [{
    Convert TF Quant ops to MHLO quant ops.
  }];

  let constructor = "mlir::quant::stablehlo::CreateConvertTFQuantOpsToMHLOPass()";
  let dependentDialects = ["TF::TensorFlowDialect", "chlo::ChloDialect",
                           "mhlo::MhloDialect", "tf_type::TFTypeDialect",
                           "quant::QuantizationDialect"];
}

using mlir::OperationState;
using mlir::UnknownLoc;
using mlir::chlo::ChloDialect;
using mlir::TF::TensorFlowDialect;
using tensorflow::tf2xla::internal::IsInBridgeAcceptableDialects;

class SharedUtilsTest : public ::testing::Test {};

TEST_F(SharedUtilsTest, IsInFunctionalDialectPasses) {
  MLIRContext context;
  context.loadDialect<TensorFlowDialect>();
  OpBuilder opBuilder(&context);

class MockSuccessPass
    : public PassWrapper<MockSuccessPass, OperationPass<ModuleOp>> {
  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<TF::TensorFlowDialect>();
  }

 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(MockSuccessPass)

  explicit MockSuccessPass() = default;

 private:
  void runOnOperation() override {

    default_device_ = std::string(default_device);
  }

  void runOnOperation() override {
    Builder builder(&getContext());
    Dialect* tf = getContext().getLoadedDialect<TensorFlowDialect>();
    getOperation().walk([&](Operation* op) {
      if (auto device_attr = op->getAttrOfType<StringAttr>(kDeviceAttr)) {
        // We assign default device to ops with device attribute that is empty.

  func.return %0 : tensor<i32>
}
)mlir";
  mlir::MLIRContext context;
  context.loadDialect<mlir::func::FuncDialect, mlir::TF::TensorFlowDialect>();
  mlir::OwningOpRef<mlir::ModuleOp> module =
      mlir::parseSourceString<mlir::ModuleOp>(code, &context);
  ASSERT_TRUE(module);
  EXPECT_TRUE(IsSupportedByReplicatedBridge(*module));
}