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

 public:
  void runOnOperation() override;
};

// A set of values that identifies a resource.
struct ResourceKey {
  StringRef device;
  StringRef container;
  StringRef shared_name;
};

ResourceKey GetResourceKey(TF::VarHandleOp var_handle_op) {
  ResourceKey resource_key;

  if (auto attr = var_handle_op->getAttrOfType<StringAttr>("device")) {

#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.

namespace mlir {
namespace TFR {

class RewriteQuantizedIOPass
    : public PassWrapper<RewriteQuantizedIOPass, OperationPass<ModuleOp>> {
 public:
  StringRef getArgument() const final { return "tfr-rewrite-quantized-io"; }

  StringRef getDescription() const final {
    return "Replaces operands and results that has quantized type with their "
           "storage types.";
  }
  void runOnOperation() override;
};

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

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

#include <string>

#include "llvm/ADT/ArrayRef.h"
#include "llvm/ADT/StringRef.h"
#include "llvm/ADT/StringSet.h"
#include "mlir/Support/LogicalResult.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tfr/ir/tfr_ops.h"

namespace mlir {
namespace TFR {
namespace {

}  // namespace

class DropSavedModelSemanticsPass
    : public PassWrapper<DropSavedModelSemanticsPass, OperationPass<ModuleOp>> {
 public:
  StringRef getArgument() const final { return "drop-savedmodel-semantics"; }
  StringRef getDescription() const final {
    return "Drops all tf_saved_model attributes";
  }

  // NOTE: The implementation is mostly copypasted from

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

#include "tensorflow/compiler/mlir/tf2xla/internal/logging_hooks.h"

#include <memory>
#include <string>

#include "llvm/ADT/StringRef.h"
#include "mlir/Pass/PassManager.h"  // from @llvm-project
#include "tensorflow/compiler/mlir/tensorflow/utils/data_dumper_logger_config.h"
#include "tensorflow/core/util/debug_data_dumper.h"

namespace tensorflow {

#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