Function TEST

tensorflow/core/util/mkl_util_test.cc:26–54

namespace {

TEST(MklUtilTest, MklDnnTfShape) {
  auto cpu_engine = engine(ENGINE_CPU, 0);
  MklDnnData<float> a(&cpu_engine);

  const int N = 1, C = 2, H = 3, W = 4;
  memory::dims a_dims = {N, C, H, W};
  MklDnnShape a_dnnl_shape;
  a_dnnl_shape.SetMklTensor(true);
  // Create TF layout in NCHW.
  a_dnnl_shape.SetTfLayout(a_dims.size(), a_dims, MKL_TENSOR_FORMAT_NCHW);
  TensorShape a_tf_shape_nchw({N, C, H, W});
  TensorShape a_tf_shape_nhwc({N, H, W, C});
  TensorShape a_dnnl_tf_shape = a_dnnl_shape.GetTfShape();
  // Check that returned shape is in NCHW format.
  EXPECT_EQ(a_tf_shape_nchw, a_dnnl_tf_shape);
  EXPECT_NE(a_tf_shape_nhwc, a_dnnl_tf_shape);

  memory::dims b_dims = {N, C, H, W};
  MklDnnShape b_dnnl_shape;
  b_dnnl_shape.SetMklTensor(true);
  // Create TF layout in NHWC.
  b_dnnl_shape.SetTfLayout(b_dims.size(), b_dims, MKL_TENSOR_FORMAT_NHWC);
  TensorShape b_tf_shape_nhwc({N, H, W, C});
  TensorShape b_tf_shape_nchw({N, C, H, W});
  TensorShape b_dnnl_tf_shape = b_dnnl_shape.GetTfShape();
  // Check that returned shape is in NHWC format.
  EXPECT_EQ(b_tf_shape_nhwc, b_dnnl_tf_shape);
  EXPECT_NE(b_tf_shape_nchw, b_dnnl_tf_shape);
}

TEST(MklUtilTest, MklDnnBlockedFormatTest) {
  // Let's create 2D tensor of shape {3, 4} with 3 being innermost dimension