From 237ce08c894a37230b9ce1e6b28e3ce46b168c12 Mon Sep 17 00:00:00 2001
From: Andy Ly <lyandy@google.com>
Date: Tue, 2 Jul 2019 16:29:35 -0700
Subject: [PATCH] [Grappler] Register DataFormatDimMap on host for CPU.

PiperOrigin-RevId: 256262935
---
 tensorflow/core/kernels/data_format_ops.cc | 10 ++++++++++
 1 file changed, 10 insertions(+)

diff --git a/tensorflow/core/kernels/data_format_ops.cc b/tensorflow/core/kernels/data_format_ops.cc
index 5c3f1b57fd7..0b4241dbb93 100644
--- a/tensorflow/core/kernels/data_format_ops.cc
+++ b/tensorflow/core/kernels/data_format_ops.cc
@@ -156,6 +156,16 @@ TF_CALL_int32(REGISTER_KERNEL);
 TF_CALL_int64(REGISTER_KERNEL);
 #undef REGISTER_KERNEL
 
+#define REGISTER_KERNEL(T)                             \
+  REGISTER_KERNEL_BUILDER(Name("DataFormatDimMap")     \
+                              .Device(DEVICE_CPU)      \
+                              .Label("host")           \
+                              .TypeConstraint<T>("T"), \
+                          DataFormatDimMapOp<CPUDevice, T>);
+TF_CALL_int32(REGISTER_KERNEL);
+TF_CALL_int64(REGISTER_KERNEL);
+#undef REGISTER_KERNEL
+
 #define REGISTER_KERNEL(T)                             \
   REGISTER_KERNEL_BUILDER(Name("DataFormatVecPermute") \
                               .Device(DEVICE_CPU)      \