const Node* arg_node = node_name_index.at("_arg0");
  EXPECT_EQ(arg_node->op_def().name(), "_Arg");
  EXPECT_EQ(arg_node->attrs().FindByString("T")->type(), DT_FLOAT);

  const Node* retval_node = node_name_index.at("_retval0");
  EXPECT_EQ(retval_node->op_def().name(), "_Retval");
  EXPECT_EQ(retval_node->attrs().FindByString("T")->type(), DT_FLOAT);

  EXPECT_EQ(identity_node->num_outputs(), 1);
  EXPECT_EQ(retval_node->num_inputs(), 1);

      if (node_op == kStatefulPartitionedCallOp) {
        // Functions called by StatefulfulPartitionedCall ops with
        // _XlaMustCompile=true are compiled by XLA.
        auto attr = node->attrs().FindByString(std::string(kMustCompileAttr));
        if (attr != nullptr && attr->b() == true) {
          return true;
        }
      }
    }
    return false;
  };