class MaterializePassthroughOpPass
    : public impl::MaterializePassthroughOpBase<MaterializePassthroughOpPass> {
 public:
  void runOnOperation() override;
};

void MaterializePassthroughOpPass::runOnOperation() {
  getOperation().walk([](TF::MlirPassthroughOp op) {
    std::string module_string(op.getMlirModule());
    // Parse the module.
    auto nested_module =

  }

  StringRef getDescription() const final {
    return "Removes `tf.AssignVariableOp(tf.VarHandleOp, tf.Const)` patterns "
           "from the initializer function of type 'restore_op'.";
  }

  void runOnOperation() override;
};

// Finds and removes the `tf.AssignVariableOp(tf.VarHandleOp, tf.Const)`
// pattern. `tf.VarHandleOp` and `tf.Const` are removed unless they are used by
// other ops.

    // 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;
};

class CastBf16OpsToF32 : public RewritePattern {
 public:
  explicit CastBf16OpsToF32(MLIRContext* context)
      : RewritePattern(MatchAnyOpTypeTag(), /*benefit=*/1, context) {}

 private:

#define GEN_PASS_DEF_TPUIDENTITYPRUNINGPASS
#include "tensorflow/compiler/mlir/tensorflow/transforms/tf_passes.h.inc"

struct TPUIdentityPruning
    : public impl::TPUIdentityPruningPassBase<TPUIdentityPruning> {
  void runOnOperation() override;
};

// Collects all reachable functions (via call ops) from a given region.
SmallVector<func::FuncOp, 4> CollectReachableFunctions(Region& region) {
  llvm::SmallPtrSet<func::FuncOp, 4> reachable_funcs;

}

class AnalyzeVariablesPass
    : public impl::AnalyzeVariablesPassBase<AnalyzeVariablesPass> {
 public:
  explicit AnalyzeVariablesPass() {}
  void runOnOperation() override;
};

void AnalyzeVariablesPass::runOnOperation() {
  auto* context = &getContext();
  auto module = getOperation();
  bool legalize_to_tfl = true;

  module.walk([&](Operation* op) {

struct AddFunctionsForExportedNamesPass
    : public impl::AddFunctionsForExportedNamesPassBase<
          AddFunctionsForExportedNamesPass> {
  void runOnOperation() override;
};

// Set the (array of) exported name(s) of a (public) function to just
// contain the given string.
void SetExportedName(func::FuncOp f, StringRef name) {
  OpBuilder b(f);

// (packed tensors) with the underlying physical TPU device.
struct TPUColocateCompositeResourceOps
    : public impl::TPUColocateCompositeResourceOpsPassBase<
          TPUColocateCompositeResourceOps> {
  void runOnOperation() override;
};

// Wraps single op in `tf_device.launch` for explicit device assignment.
void WrapOpInLaunch(OpBuilder* builder, Location loc, Operation* op,
                    llvm::StringRef device) {

 public:
  MLIR_DEFINE_EXPLICIT_INTERNAL_INLINE_TYPE_ID(UnwrapXlaCallModuleOpPass)

  explicit UnwrapXlaCallModuleOpPass() = default;

 private:
  void runOnOperation() override;
};

void UnwrapXlaCallModuleOp(TF::XlaCallModuleOp call_op,
                           SymbolTable& symbol_table) {
  // Do not inline lifted quantized functions used for fusing patterns.

class InitializeVariablesInSessionInitializerPass
    : public impl::InitializeVariablesInSessionInitializerPassBase<
          InitializeVariablesInSessionInitializerPass> {
 public:
  void runOnOperation() final {
    static tensorflow::Session* session = new TF::test_util::FakeSession();
    if (failed(tf_saved_model::InitializeVariablesInSessionInitializer(
            getOperation(), session)))
      signalPassFailure();

// internal variants by introducing XlaRecvTPUEmbeddingDeduplicationData op.
struct RewriteTPUEmbeddingOps
    : public impl::RewriteTPUEmbeddingOpsPassBase<RewriteTPUEmbeddingOps> {
  void runOnOperation() override;
};

// Rewrites the given op to `OpT` op after adding the given operand at the end.
template <typename OpT>
OpT AddOperandAndRewriteAs(Operation* op, Value operand, NamedAttrList attr,