Function MklTensorFormatToMklDnnDataFormat

tensorflow/core/util/mkl_util.h:1010–1020

Map MklTensorFormat to OneDNN format tag @input: MklTensorFormat i.e. TensorFlow data format @return: OneDNN's memory format tag corresponding to MklTensorFormat. Fails with an error if invalid data format.

// @return: OneDNN's memory format tag corresponding to MklTensorFormat.
//          Fails with an error if invalid data format.
inline memory::format_tag MklTensorFormatToMklDnnDataFormat(
    MklTensorFormat format) {
  if (format == MklTensorFormat::FORMAT_NHWC) return memory::format_tag::nhwc;
  if (format == MklTensorFormat::FORMAT_NCHW) return memory::format_tag::nchw;
  if (format == MklTensorFormat::FORMAT_NDHWC) return memory::format_tag::ndhwc;
  if (format == MklTensorFormat::FORMAT_NCDHW) return memory::format_tag::ncdhw;
  if (format == MklTensorFormat::FORMAT_X) return memory::format_tag::x;
  if (format == MklTensorFormat::FORMAT_NC) return memory::format_tag::nc;
  if (format == MklTensorFormat::FORMAT_TNC) return memory::format_tag::tnc;
  return memory::format_tag::undef;
}

/// Map TensorFlow data format into OneDNN 3D data format
/// @input: TensorFlow data format