Moves MemoryType inference code out of OpKernel so that it can reused.

Change: 114448861

1 files changed, 0 insertions, 37 deletions

diff --git a/tensorflow/core/framework/op_kernel_test.cc b/tensorflow/core/framework/op_kernel_test.cc
index a531b10cde..4275ad0f73 100644
--- a/tensorflow/core/framework/op_kernel_test.cc
+++ b/tensorflow/core/framework/op_kernel_test.cc
@@ -745,43 +745,6 @@ TEST_F(GetAttrTest, TypeList) {
   EXPECT_EQ(DT_BOOL, get_attr_kernel->type_vector[1]);
 }
 
-REGISTER_OP("HostMemoryTest")
-    .Input("a: float")
-    .Input("b: T")
-    .Input("c: N * string")
-    .Output("o: N * T")
-    .Attr("T: type")
-    .Attr("N: int");
-REGISTER_KERNEL_BUILDER(Name("HostMemoryTest").Device(DEVICE_CPU), DummyKernel);
-REGISTER_KERNEL_BUILDER(Name("HostMemoryTest")
-                            .Device(DEVICE_GPU)
-                            .HostMemory("a")
-                            .HostMemory("c")
-                            .HostMemory("o"),
-                        DummyKernel);
-
-TEST(MemoryTypesForNode, Simple) {
-  NodeDef node_def;
-  TF_ASSERT_OK(NodeDefBuilder("test", "HostMemoryTest")
-                   .Input(FakeInput())
-                   .Input(FakeInput(DT_BOOL))
-                   .Input(FakeInput(3))
-                   .Finalize(&node_def));
-  MemoryTypeVector input, output;
-
-  TF_EXPECT_OK(MemoryTypesForNode(*OpRegistry::Global(), DEVICE_CPU, node_def,
-                                  &input, &output));
-  EXPECT_EQ(MemoryTypeVector(5, DEVICE_MEMORY), input);
-  EXPECT_EQ(MemoryTypeVector(3, DEVICE_MEMORY), output);
-
-  TF_EXPECT_OK(MemoryTypesForNode(*OpRegistry::Global(), DEVICE_GPU, node_def,
-                                  &input, &output));
-  EXPECT_EQ(MemoryTypeVector({HOST_MEMORY, DEVICE_MEMORY, HOST_MEMORY,
-                              HOST_MEMORY, HOST_MEMORY}),
-            input);
-  EXPECT_EQ(MemoryTypeVector(3, HOST_MEMORY), output);
-}
-
 class BaseKernel : public ::tensorflow::OpKernel {
  public:
   explicit BaseKernel(OpKernelConstruction* context) : OpKernel(context) {}