// and input devices, variables used by this function.
    llvm::SmallVector<llvm::StringRef, 4> input_names;
    llvm::SmallVector<
        std::pair<tensorflow::DataType, tensorflow::PartialTensorShape>, 4>
        input_specs;
    llvm::SmallVector<llvm::StringRef, 4> input_devices;
    llvm::SmallVector<mlir::Operation*, 4> bound_inputs;

    ModuleOp module, Session* session,
    const SmallSet<StringRef, 4>& resource_names) {
  OpBuilder builder(module.getBodyRegion());

  if (!session) return module.emitOpError() << "no session provided";

  // Read all resource variables from the session.
  std::vector<std::string> variable_names;
  variable_names.reserve(resource_names.size());
  for (StringRef name : resource_names) variable_names.push_back(name.str());

#ifndef TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_LEGALIZE_TF_MLIR_H_
#define TENSORFLOW_COMPILER_MLIR_TF2XLA_INTERNAL_LEGALIZE_TF_MLIR_H_

#include <string>
#include <vector>

#include "llvm/ADT/StringRef.h"
#include "mlir/Pass/Pass.h"  // from @llvm-project
#include "tensorflow/compiler/tf2xla/xla_helpers.h"
#include "tensorflow/core/tpu/kernels/tpu_compile_op_support.h"
#include "tsl/platform/statusor.h"

// to represent calling a TF host callback function.
constexpr llvm::StringRef kTfTargetName = "tf.call_tf_function";

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

}  // namespace

#ifndef TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_UTILS_TF_TYPE_UTILS_H_
#define TENSORFLOW_COMPILER_MLIR_QUANTIZATION_STABLEHLO_UTILS_TF_TYPE_UTILS_H_

#include "llvm/ADT/StringRef.h"
#include "mlir/IR/BuiltinAttributes.h"  // from @llvm-project
#include "mlir/IR/BuiltinTypes.h"  // from @llvm-project
#include "mlir/IR/Types.h"  // from @llvm-project

==============================================================================*/

#include "tensorflow/compiler/mlir/tensorflow/utils/mlprogram_util.h"

#include "llvm/ADT/STLFunctionalExtras.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/Twine.h"
#include "mlir/IR/BuiltinOps.h"  // from @llvm-project
#include "mlir/Pass/PassManager.h"  // from @llvm-project
#include "mlir/Pass/PassRegistry.h"  // from @llvm-project

#include "xla/xla_data.pb.h"
#include "tensorflow/core/protobuf/tpu/compile_metadata.pb.h"

namespace tensorflow {

inline constexpr llvm::StringRef kInputShardingAttr =
    "input_sharding_configuration";
inline constexpr llvm::StringRef kOutputShardingAttr =
    "output_sharding_configuration";

// Parses the sharding string. This sharding string can be binary (serialized)
// or human readable.

#include <utility>

#include "absl/status/status.h"
#include "absl/strings/string_view.h"
#include "llvm/ADT/DenseMap.h"
#include "llvm/ADT/STLExtras.h"
#include "llvm/ADT/SmallVector.h"
#include "llvm/ADT/StringRef.h"
#include "mlir/Dialect/Func/IR/FuncOps.h"  // from @llvm-project
#include "mlir/IR/Attributes.h"  // from @llvm-project
#include "mlir/IR/Builders.h"  // from @llvm-project

namespace mlir {
namespace quant {
namespace {

using ::mlir::tf_saved_model::GetInitializerFunction;
using ::mlir::tf_saved_model::kTfSavedModelInitializerRestoreType;

constexpr StringRef kTfQuantSaveV2OpName = "tf_quant__save_save_v2";
constexpr StringRef kTfQuantSaveReturnOpName = "tf_quant__save_return";

// A pass that creates a new function that wraps the newly created SaveV2 op.

#ifndef TENSORFLOW_COMPILER_MLIR_TENSORFLOW_TRANSLATE_EXPORT_TF_DIALECT_OP_H_
#define TENSORFLOW_COMPILER_MLIR_TENSORFLOW_TRANSLATE_EXPORT_TF_DIALECT_OP_H_

#include "llvm/ADT/StringRef.h"
#include "mlir/IR/Operation.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/utils/export_utils.h"
#include "tensorflow/core/framework/attr_value.pb.h"
#include "tensorflow/core/framework/node_def.pb.h"