// devices, and empty StringRef will be returned. Control dependencies,
// NextIteration.Source -> NextIteration.Sink token dependencies, and
// LoopCond -> Switch data dependencies are ignored.
llvm::StringRef FindDeviceFromOperands(
    Operation& op,
    const llvm::DenseMap<Value, llvm::StringRef>& value_to_device) {
  llvm::StringRef new_device;

#include "tensorflow/compiler/mlir/tensorflow/transforms/passes.h"

namespace mlir {
namespace TF {
namespace {

// Location attribute.
constexpr StringRef kClassAttr = "_class";
constexpr StringRef kSharedNameAttr = "shared_name";
constexpr StringRef kLocationPrefix = "loc:@";

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

// function.
constexpr llvm::StringRef kStablehloEntryFunctionAttrName = "_entry_function";

// Every stablehlo function deserialized from XlaCallModule has this attribute.
constexpr llvm::StringRef kFromXlaCallModuleAttrName = "_from_xla_call_module";

// Name of `tf.XlaCallModule`'s dictionary attribute for keeping the
// deserialized stablehlo module's attributes.
constexpr llvm::StringRef kStablehloModuleAttrsAttrName =

// TF_DUMP_GRAPH_PREFIX environment variable if `dirname` is empty and
// suffixing `name` with ".mlir".
std::string DumpRawStringToFile(llvm::StringRef name, llvm::StringRef content,
                                llvm::StringRef dirname = "");

// Enable the crash reproducer on the provided PassManager to the provided
// directory path.
// If the provided path is empty, it is retrieved from the

// CHECK-LABEL: func @testTFDialect
func.func @testTFDialect(%arg0: tensor<4x?x!tf_type.stringref>) -> tensor<4x2x!tf_type.string> {
  %0 = "tf.Identity"(%arg0) : (tensor<4x?x!tf_type.stringref>) -> tensor<4x2x!tf_type.string>
  func.return %0 : tensor<4x2x!tf_type.string>
}


// -----

func.func @testTFDialect(%arg0: tensor<4x?x!tf_type.stringref>) -> tensor<4x2x!tf_type.string> {

    mlir::ModuleOp module_op, llvm::StringRef device_type,
    xla::XlaComputation* xla_computation, bool use_tuple_args,
    bool enable_op_fallback, bool return_tuple,
    const XlaShapeLayoutHelpers::ShapeDeterminationFns shape_determination_fns =
        {},
    llvm::MutableArrayRef<std::unique_ptr<mlir::Pass>>
        custom_legalization_passes = {},
    llvm::StringRef module_name = llvm::StringRef());

bool is_valid_special_tpu_op(
    std::vector<IslandOp>& ops, llvm::StringRef cluster_name,
    llvm::SmallDenseMap<llvm::StringRef, llvm::SmallDenseSet<Operation*>>&
        cluster_to_tpu_op_map) {
  for (IslandOp op : ops) {
    Operation* wrapped_op = &op.GetBody().front();
    std::optional<llvm::StringRef> wrapped_op_cluster_name =
        GetTpuClusterName(wrapped_op);

    ctx_.loadDialect<TensorFlowSavedModelDialect, func::FuncDialect>();
  }

  MLIRContext ctx_;
};

// Parses `module_op_str` and returns the resulting `ModuleOp`.
ModuleOp ParseModuleOp(const StringRef module_op_str, Block& block,
                       MLIRContext& ctx) {
  const LogicalResult parse_result =
      parseSourceString(module_op_str, &block, ParserConfig(&ctx));
  EXPECT_TRUE(succeeded(parse_result));

namespace mlir {

class MLIRContext;

namespace TF {

SmallVector<int64_t, 4> ReversePermutation(ArrayRef<int64_t> permutation);

SmallVector<int64_t, 4> GetDataFormatPermutation(StringRef from, StringRef to);

// Shuffle elements in the `attr` according to the permutation. Optional
// `inner_size` allows to shuffle array attributes created from rank 2 tensors
// on outer dimension only.

    // Op type is a NameLoc with the ":" suffix.
    StringRef op_type_with_suffix =
        mlir::cast<NameLoc>(locations.front()).getName().strref();
    StringRef op_type =
        op_type_with_suffix.substr(0, op_type_with_suffix.size() - 1);
    new_unit.set_op_type(op_type.str());

    if (isa<NameLoc>(locations.back())) {
      StringRef name_loc_id =