ASSERT_EQ(errors::OK, status_.code()) << status_.message();
    x2.reset(x2_raw);
  }

  status_ = registry_.Register("IdentityN", IdentityNRegisterer);
  ASSERT_EQ(errors::OK, status_.code()) << status_.message();
  auto IdentityNGradModel = BuildGradModel(IdentityNModel, registry_);

  std::vector<AbstractTensorHandle*> outputs(2);
  status_ =
      RunModel(IdentityNGradModel, immediate_execution_ctx_.get(),

  context->appendDialectRegistry(registry_);
  context->loadAllAvailableDialects();

  // Run TAC passes.
  auto status = RunTacPasses(&module, options_.debug_mode);

  if (!status.ok()) {
    return status;
  }

  return exporter_->Export(module);
}

void TacModule::RegisterExtraDialects(mlir::DialectRegistry& registry) {
  registry.appendTo(registry_);
}
}  // namespace tac

  /// the gradient function for 'op' exists in the registry.
  /// Note that 'func' can be null for ops that have registered no-gradient with
  /// the registry.
  /// Returns error status otherwise.
  Status Lookup(const string& op, GradFunc* func) const;

  /// Returns a pointer to the global gradient function registry.
  static GradOpRegistry* Global();

 private:

namespace tensorflow {
namespace gradients {
TapeOperation::TapeOperation(AbstractOperation* parent_op, Tape* tape,
                             const GradientRegistry& registry)
    : AbstractOperation(kTape),
      parent_op_(parent_op),
      tape_(tape),
      registry_(registry) {
  // TODO(b/172003047): Consider making AbstractOperation RefCounted.
  // parent_op_->Ref();
}
void TapeOperation::Release() {

  const std::vector<const tac::TargetHardware*>& GetAvailableHardwares() const {
    return const_backends_;
  }

  // Registers all dialects in 'registry' with the module.
  // This to allow clients to register extra dialects required.
  void RegisterExtraDialects(mlir::DialectRegistry& registry);

 protected:
  // Adds TAC passes to the 'pass_manager'.
  virtual void AddTACPass(mlir::OpPassManager* pass_manager,

//   // More complex gradient functions can use inputs/attrs etc. from the
//   // forward `op`.
//   return new AddGradientFunction;
// }
//
// Status RegisterGradients(GradientRegistry* registry) {
//   return registry->Register("Add", AddRegisterer);
// }
class GradientFunction {
 public:
  virtual Status Compute(AbstractContext* ctx,
                         absl::Span<AbstractTensorHandle* const> grad_outputs,

#include "tensorflow/compiler/mlir/tfr/ir/tfr_ops.h"

int main(int argc, char **argv) {
  tensorflow::InitMlir y(&argc, &argv);

  mlir::registerAllPasses();

  mlir::DialectRegistry registry;
  registry.insert<mlir::scf::SCFDialect, mlir::TF::TensorFlowDialect,
                  mlir::arith::ArithDialect, mlir::func::FuncDialect,
                  mlir::shape::ShapeDialect, mlir::quant::QuantizationDialect,

                  mlir::tf_executor::TensorFlowExecutorDialect,
                  mlir::vhlo::VhloDialect>();
  mlir::mhlo::registerAllMhloDialects(registry);
  mlir::func::registerAllExtensions(registry);
  return failed(
      mlir::MlirOptMain(argc, argv, "StableHLO quant Pass Driver\n", registry));

int main(int argc, char **argv) {
  tensorflow::InitMlir y(&argc, &argv);

  mlir::registerAllPasses();
  mlir::registerTensorFlowPasses();

  mlir::DialectRegistry registry;
  registry.insert<mlir::scf::SCFDialect, mlir::TF::TensorFlowDialect,
                  mlir::tf_saved_model::TensorFlowSavedModelDialect,
                  mlir::func::FuncDialect, mlir::shape::ShapeDialect,

         "or compilation failure from the MLIR bridge (full fallback mode).";

  mlir::DialectRegistry registry;
  mlir::RegisterAllTensorFlowDialects(registry);
  mlir::mhlo::registerAllMhloDialects(registry);
  mlir::stablehlo::registerAllDialects(registry);
  mlir::MLIRContext context(registry);
  mlir::OwningOpRef<mlir::ModuleOp> mlir_module;
  TF_RETURN_IF_ERROR(