Function MklDnnDimsInNCHW

tensorflow/core/util/mkl_util.h:1122–1135

Overloaded version of function TFShapeToMklDnnDimsInNCHW above. Input parameters are self-explanatory.

/// Overloaded version of function TFShapeToMklDnnDimsInNCHW above.
/// Input parameters are self-explanatory.
inline memory::dims MklDnnDimsInNCHW(const memory::dims& in_dims,
                                     TensorFormat format) {
  // Validate format.
  DCHECK_NE(TFDataFormatToMklDnnDataFormat(format),
            MklTensorFormat::FORMAT_INVALID);

  int n = in_dims[GetTensorDimIndex(format, 'N')];
  int c = in_dims[GetTensorDimIndex(format, 'C')];
  int h = in_dims[GetTensorDimIndex(format, 'H')];
  int w = in_dims[GetTensorDimIndex(format, 'W')];

  // OneDNN requires dimensions in NCHW format.
  return memory::dims({n, c, h, w});
}

/// Overloaded version of function TFShapeToMklDnnDimsInNCDHW above.
/// Input parameters are self-explanatory.