namespace {

class BreakUpIslands : public TF::PerFunctionAggregateAnalysisConsumerPass<
                           BreakUpIslands, TF::SideEffectAnalysis> {
  void getDependentDialects(DialectRegistry& registry) const override {
    registry.insert<tf_executor::TensorFlowExecutorDialect>();
  }

 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(BreakUpIslands)

  //
  //  Arg  --> B  --> C  --> D --> Retval
  class Subgraph {
   public:
    // Creates a graph to build the subgraph in, if it doesn't already exist,
    // using the same op registry and versions as graph_in.
    Node* MakeNodeImage(const Graph* graph_in, Node* node);

    // Returns the graph the subgraph is being built in.
    Graph* GetGraph() const;

  let description = [{
    Fixed point shape refinement pass that utilizes the shape functions
    registered on ops using the InferTypeOpInterface as well as by bridging to
    the TensorFlow op registry's shape functions. This is an interprocedural
    pass that propagates information across function calls/control flow
    operations where possible (the GuaranteeAllFuncsOneUsePass is often run

}

void LoadImporterDialects(mlir::MLIRContext& context) {
  // Load dialects involved in the conversion
  mlir::DialectRegistry registry;
  mlir::RegisterAllTensorFlowDialectsImpl(registry, false);
  context.appendDialectRegistry(registry);
  for (llvm::StringRef name : registry.getDialectNames())
    context.getOrLoadDialect(name);
}

absl::StatusOr<std::string> GetDenseTensorNameFromTensorInfo(

    }
  }
  return false;
}

struct EmbeddingPipeliningPass
    : public ::impl::EmbeddingPipeliningPassBase<EmbeddingPipeliningPass> {
  void getDependentDialects(mlir::DialectRegistry& registry) const override {
    registry.insert<TF::TensorFlowDialect>();
  }

  void runOnOperation() override;
};

bool UseEmbeddingPipelining(ModuleOp& module) {
  // Enable automated pipelining pass unless:

        *registry, InferSparseSegmentOpCost<mlir::TF::SparseSegmentMeanOp>);
    RegisterCostFunction<mlir::TF::SparseSegmentSqrtNOp>(
        *registry, InferSparseSegmentOpCost<mlir::TF::SparseSegmentSqrtNOp>);
    RegisterCostFunction<mlir::TF::LookupTableFindV2Op>(
        *registry, InferLookupTableFindV2Cost);
    return registry;
  }();
  return *registry;
}

}  // namespace

                                        TF_Status* status) {
  RegisterPasses();
  mlir::DialectRegistry registry;
  mlir::RegisterAllTensorFlowDialects(registry);
  mlir::stablehlo::registerAllDialects(registry);
  registry.insert<mlir::shape::ShapeDialect>();
  mlir::MLIRContext context(registry);
  mlir::OwningOpRef<mlir::ModuleOp> module;
  {
    mlir::StatusScopedDiagnosticHandler diagnostic_handler(&context);

mlir::DialectRegistry GetDefaultDialectRegistry() {
  mlir::DialectRegistry registry;

  registry.insert<arith::ArithDialect>();
  registry.insert<func::FuncDialect>();
  registry.insert<tensor::TensorDialect>();
  registry.insert<shape::ShapeDialect>();
  registry.insert<TF::TensorFlowDialect>();
  registry.insert<chlo::ChloDialect>();

  return registry;
}

class XlaLegalizeTargetsTest : public testing::Test {
 public:

// Creates an MLIRContext with the extensions required for quantization are
// registered.
inline std::unique_ptr<MLIRContext> CreateMlirContextForQuantization() {
  DialectRegistry registry{};
  func::registerAllExtensions(registry);
  return std::make_unique<MLIRContext>(registry);
}

}  // namespace mlir::quant::stablehlo

### TPU Enhancements

*   Adds 3D mesh support in TPU configurations ops.
*   Added TPU code for `FTRL` with `multiply_linear_by_lr`.
*   Silently adds a new file system registry at `gstpu`.
*   Support `restartType` in cloud tpu client.
*   Depend on a specific version of google-api-python-client.
*   Fixes apiclient import.

### Tracing and Debugging