Function TEST

tensorflow/compiler/jit/async_io_conversion_pass_test.cc:163–208

Tests that AsyncIoConversionPass does not replace either Resource or Const output edges.

// Tests that AsyncIoConversionPass does not replace either Resource
// or Const output edges.
TEST(AsyncIoConversionTest, ResourceAndConstRetval) {
  const string cluster_name = "cluster_var_and_const";

  auto graph = absl::make_unique<Graph>(OpRegistry::Global());
  FunctionDefLibrary fdef_lib;
  FunctionLibraryDefinition flib_def(OpRegistry::Global(), fdef_lib);
  {
    GraphDefBuilder builder(GraphDefBuilder::kFailImmediately, &flib_def);
    Node* res_in =
        ops::SourceOp("VarHandleOp", builder.opts()
                                         .WithName("varhandle")
                                         .WithAttr("dtype", DT_FLOAT)
                                         .WithAttr("shape", TensorShape({})));
    Node* const_in =
        ops::SourceOp("Const", builder.opts()
                                   .WithName("")
                                   .WithAttr("dtype", DT_FLOAT)
                                   .WithAttr("value", Tensor())
                                   .WithAttr(kXlaClusterAttr, cluster_name));
    Node* identity0 = ops::UnaryOp(
        "Identity", res_in, builder.opts()
                                .WithName("identity0_in_cluster")
                                .WithDevice(kTestDeviceName)
                                .WithAttr("T", DT_RESOURCE)
                                .WithAttr(kXlaClusterAttr, cluster_name));
    Node* identity1 = ops::UnaryOp("Identity", identity0,
                                   builder.opts()
                                       .WithName("identity1_outside_cluster")
                                       .WithDevice(kTestDeviceName)
                                       .WithAttr("T", DT_RESOURCE));
    Node* identity2 = ops::UnaryOp("Identity", const_in,
                                   builder.opts()
                                       .WithName("identity2_outside_cluster")
                                       .WithDevice(kTestDeviceName)
                                       .WithAttr("T", DT_FLOAT));

    TF_EXPECT_OK(GraphDefBuilderToGraph(builder, graph.get()));
  }

  TF_ASSERT_OK(ConvertToAsyncOut(&graph, &flib_def));

  // Verify that no AsyncOut is inserted.
  for (const Node* n : graph->op_nodes()) {
    EXPECT_TRUE(!absl::StartsWith(n->type_string(), "_XlaAsyncOut"));
  }
}

}  // namespace
}  // namespace tensorflow