test_mode_ = other.test_mode_;
    op_set_ = other.op_set_;
    initializeForTest();
  }

  StringRef getArgument() const final {
    // This is the argument used to refer to the pass in
    // the textual format (on the commandline for example).
    return "quant-lift-quantizable-spots-as-functions";
  }

  StringRef getDescription() const final {
    // This is a brief description of the pass.

limitations under the License.
==============================================================================*/

#include <memory>
#include <optional>
#include <string>
#include <utility>

#include "llvm/ADT/StringRef.h"
#include "llvm/Support/Debug.h"
#include "llvm/Support/FormatVariadic.h"
#include "mlir/Dialect/Arith/IR/Arith.h"  // from @llvm-project
#include "mlir/Dialect/Func/Extensions/AllExtensions.h"  // from @llvm-project

namespace mlir {
namespace TFL {
namespace {
#define GEN_PASS_DEF_LIFTTFLITEFLEXOPSPASS
#include "tensorflow/compiler/mlir/lite/transforms/passes.h.inc"

using ::tensorflow::StatusOr;

constexpr StringRef kFlexOpNamePrefix = "Flex";

// Pattern that converts TFL::CustomOp that encodes a Flex op into a TF dialect
// operation.
class LiftFlexCustomOp : public OpRewritePattern<TFL::CustomOp> {
 public:

// `tf.backend_config` is a DictionaryAttr, JAX2TF sets the value of its
// i64 attribute `called_index` to the TF function's name.
constexpr llvm::StringRef kTfBackendConfigAttrName = "tf.backend_config";
constexpr llvm::StringRef kCalledIndexAttrName = "called_index";
constexpr llvm::StringRef kCalledFuncAttrName = "called_func";

// Deserialize the StableHLO module embedded in XlaCallModuleOp's module
// attribute.

    absl::flat_hash_set<Node*>* control_ret_nodes) {
  mlir::Block& block = function.front();

  // Extract input & output names if set.
  llvm::SmallVector<llvm::StringRef, 2> input_names;
  llvm::SmallVector<llvm::StringRef, 2> output_names;
  llvm::SmallVector<llvm::StringRef, 2> unique_output_names;
  auto dict_attr =
      function->getAttrOfType<mlir::DictionaryAttr>(kEntryFuncAttr);
  if (dict_attr) {

    absl::flat_hash_set<Node*>* control_ret_nodes) {
  mlir::Block& block = function.front();

  // Extract input & output names if set.
  llvm::SmallVector<llvm::StringRef, 2> input_names;
  llvm::SmallVector<llvm::StringRef, 2> output_names;
  llvm::SmallVector<llvm::StringRef, 2> unique_output_names;
  auto dict_attr =
      function->getAttrOfType<mlir::DictionaryAttr>(kEntryFuncAttr);
  if (dict_attr) {

std::optional<int64_t> GetDotGeneralQuantizationDim(
    ::mlir::stablehlo::DotGeneralOp dot_general_op);

// Checks if a `StringRef` contains 'conv' or 'dot_general'.
bool ContainsConvOrDot(StringRef str);

}  // namespace mlir::quant

// handled differently as individual names are used per data send and receive.
void SetFrontendAttributes(Operation* op, int32_t index, StringRef key,
                           Type type, bool device_to_host,
                           StringRef host_handler_name) {
  MLIRContext* context = op->getContext();

  std::string formatted_key =
      device_to_host ? llvm::formatv("{0}_dtoh_{1}", key, index).str()

#include "tensorflow/core/platform/protobuf.h"
#include "tensorflow/core/util/tensor_bundle/byte_swap_tensor.h"

namespace tensorflow {

static absl::StatusOr<mlir::OwningOpRef<mlir::ModuleOp>> GraphdefToMlirImport(
    llvm::StringRef input, const std::vector<std::string>& input_arrays,
    const std::vector<std::string>& input_dtypes,
    const std::vector<std::optional<std::vector<int>>>& input_shapes,
    const std::vector<std::string>& output_arrays,

  auto string_values = string_attr.getRawStringData();
  ASSERT_EQ(string_values.size(), 4);
  EXPECT_EQ(string_values[0], mlir::StringRef("one"));
  EXPECT_EQ(string_values[1], mlir::StringRef("two"));
  EXPECT_EQ(string_values[2], mlir::StringRef("three"));
  EXPECT_EQ(string_values[3], mlir::StringRef("four"));
}

class ConvertTensorTest : public ::testing::Test {
 protected:
  template <typename T>