void ConvertLSTMCellSimpleToFusedLSTM::UpdateFuncSignature() {
  // https://github.com/tensorflow/community/pull/113
  SmallVector<int64_t, 2> output_shape{1, tensorflow::kTFDynamicSize};
  auto input_types = fused_func_op_.getFunctionType().getInputs();
  auto output_type = tensorflow::GetTypeFromTFTensorShape(
      output_shape,
      mlir::cast<RankedTensorType>(input_.getType()).getElementType());

  OpBuilder builder(module_op.getContext());
  builder.setInsertionPointAfter(session_init_op);

  const Location loc = builder.getUnknownLoc();
  const auto func_type = builder.getFunctionType(/*inputs=*/{}, /*results=*/{});

  auto init_func = builder.create<func::FuncOp>(
      loc, /*sym_name=*/"init_func_restore_op", func_type);
  builder.createBlock(&init_func.getBody(), /*insertPt=*/init_func.begin(),

      }
      // The FuncOp type can contain types that the op's operand and result
      // types do not contain.
      if (auto func = dyn_cast<func::FuncOp>(op)) {
        if (!converter_.isSignatureLegal(func.getFunctionType())) return false;
      }
      return converter_.isLegal(op);
    });
  }

 private:
  TFQuantTypeConverter &converter_;
};

class TFQuantTypePattern : public ConversionPattern {

      return WalkResult::interrupt();
    }

    tensorflow::IncreaseOpExpansionExecuteCounterByOne(
        op->getName().getStringRef().str());

    auto compose_func_type = compose_func.getFunctionType();
    builder.setInsertionPoint(op);
    TFRTensorType unconstrainted_tensor_type = builder.getType<TFRTensorType>();

    // Create the new operands. This is mapping the operands from the target

      terminator->setOperand(i, dequantized_result);
      returned_op->erase();
    } else {
      output_types.push_back(returned_value.getType());
    }
  }
  auto new_func_type = builder.getFunctionType(input_types, output_types);
  func.setType(new_func_type);
}

enum RemoveVolatileOpsType {
  // Remove all volatile quant-dequant ops.
  kPreserveNone,

        mlir::cast<FlatSymbolRefAttr>(sym_ref).getValue());

    if (!init_func_op)
      return session_initializer.emitOpError()
             << "the initializer function does not exist";

    if (!init_func_op.getFunctionType().getResults().empty())
      return session_initializer.emitOpError()
             << "the initializer function should have no output";

    auto exported_names = GetExportedNames(init_func_op);

    // TODO(aminim) figure out the location story here
    module_ = ModuleOp::create(builder_.getUnknownLoc());
    func_ = func::FuncOp::create(
        builder_.getUnknownLoc(), name,
        builder_.getFunctionType(std::nullopt, std::nullopt));
    module_->push_back(func_);
    builder_ = OpBuilder::atBlockBegin(func_.addEntryBlock());
  }

  void Release() override { delete this; }

  }

  SmallVector<Type> output_types;
  for (const Value output : xla_call_module_op.getOutput()) {
    output_types.push_back(output.getType());
  }

  entry_func_op.setFunctionType(
      rewriter.getFunctionType(arg_types, output_types));

  // Replace argument types and locs.
  Block& entry = entry_func_op->getRegion(0).front();
  for (auto [arg, arg_type, arg_loc] :

  builder.setInsertionPointToEnd(&module.getBodyRegion().back());
  func::FuncOp func_op = builder.create<func::FuncOp>(
      module.getLoc(), name,
      builder.getFunctionType(input_types, output_types));
  func_op.setPrivate();

  return func_op;
}

TF::StatefulPartitionedCallOp EncapsulateOpsInFunc(
    OpBuilder& builder, const llvm::SetVector<Operation*>& ops,