tensor.flat<float>()(i) = offset + i;
  }
  return tensor;
}

NodeDef MakeNode(
    absl::string_view name, absl::string_view op,
    absl::Span<const std::string> inputs,
    absl::Span<
        const std::pair<std::string, FunctionDefHelper::AttrValueWrapper>>
        attrs) {
  NodeDef node;
  node.set_name(std::string(name));
  node.set_op(std::string(op));

namespace mlir::quant::stablehlo {
namespace {

using ::tensorflow::AssetFileDef;
using ::tensorflow::GraphDef;
using ::tensorflow::NodeDef;
using ::tensorflow::SaverDef;
using ::tensorflow::quantization::ExportedModel;
using ::testing::HasSubstr;
using ::testing::IsEmpty;
using ::testing::SizeIs;
using ::testing::StrEq;
using ::tsl::protobuf::TextFormat;
using ::tsl::testing::IsOk;
using ::tsl::testing::StatusIs;

  absl::StatusOr<std::unique_ptr<tensorflow::NodeDef>> node_def =
      tensorflow::ConvertTFDialectOpToNodeDef(
          op, /*name=*/"", /*ignore_unregistered_attrs=*/true);
  if (!node_def.ok()) {
    op->emitError("Failed to obtain TensorFlow NodeDef: " +
                  node_def.status().ToString());
    return false;
  }
  std::string node_def_str;
  if (!(*node_def)->SerializeToString(&node_def_str)) {

  Status Run(EagerOperation* orig_op,
             std::unique_ptr<tensorflow::EagerOperation>* out_op) override;

 private:
  // Whether to run this pass. If this is enabled, the NodeDef will be imported
  // to MLIR even no tf composition file is found.
  bool IsEnabled() {
    const char* tfr_lib_env_val = getenv(string(kTFRLibEnv).c_str());
    return tfr_lib_env_val != nullptr;
  }
};

    const FunctionDef *true_def = fld.Find("oc_then_branch_host_if_true_fn");
    EXPECT_NE(true_def, nullptr);
    bool has_identity_true_fn_node = false;
    for (const auto &node_def : true_def->node_def()) {
      if (node_def.name() == "identity_true_fn") {
        has_identity_true_fn_node = true;
        break;
      }
    }
    EXPECT_TRUE(has_identity_true_fn_node);

                                       std::vector<int>* resource_arg_indices);

// Given a NodeDef `node_def` returns true iff `node_def` has kXlaCompileAttr
// set.
bool CanCreateXlaKernel(const NodeDef& node_def);

// Returns memory types for the input.
// `constant_arg_indices` and `resource_arg_indices` are sorted arrays of

// runtime OpKernel.
void RemoveIdentityCast(NodeDef* node_def) {
  auto attr = node_def->mutable_attr();
  if (node_def->op() == "Cast" && attr->contains("SrcT") &&
      attr->contains("DstT") &&
      attr->at("SrcT").type() == attr->at("DstT").type() &&
      attr->contains("Truncate") && !attr->at("Truncate").b()) {
    node_def->set_op("Identity");
    attr->insert({{"T", attr->at("SrcT")}});

  // Since we may need to load many constants, we create a hash map of these
  // names to their corresponding nodes at load time in order to look them up
  // in constant time.
  tensorflow::StatusOr<const tensorflow::NodeDef*> GetGraphDefNode(
      std::string name);

  // Returns a list of function defs in the SavedModel.
  const protobuf::RepeatedPtrField<FunctionDef>& GetFunctionDefs();

  kUncompiled,
  kCompiling,
  kCompiled,
};

// Creates a single-node graph using the specified `node_def` as the only op
// apart from the arg and retval nodes corresponding to `args` and
// `result_types` respectively.
absl::StatusOr<std::unique_ptr<Graph>> CreateSingleOpGraph(
    const NodeDef& node_def, absl::Span<const XlaArgument> args,
    absl::Span<const DataType> result_types);

absl::StatusOr<FunctionDef> TFRDecomposeContext::ExpandNode(
    const NodeDef& node_def, StringPiece func_name) {
  const OpDef* op_def;
  TF_RETURN_IF_ERROR(OpRegistry::Global()->LookUpOpDef(node_def.op(), &op_def));
  DataTypeVector input_dtys, output_dtys;
  TF_RETURN_IF_ERROR(InputTypesForNode(node_def, *op_def, &input_dtys));
  TF_RETURN_IF_ERROR(OutputTypesForNode(node_def, *op_def, &output_dtys));