Overloaded version of function TFShapeToMklDnnDimsInNCDHW above. Input parameters are self-explanatory.
| 1137 | /// Overloaded version of function TFShapeToMklDnnDimsInNCDHW above. |
| 1138 | /// Input parameters are self-explanatory. |
| 1139 | inline memory::dims MklDnnDimsInNCDHW(const memory::dims& in_dims, |
| 1140 | TensorFormat format) { |
| 1141 | // Validate format. |
| 1142 | DCHECK_NE(TFDataFormatToMklDnnDataFormat(format), |
| 1143 | MklTensorFormat::FORMAT_INVALID); |
| 1144 | |
| 1145 | int n = in_dims[GetTensorDimIndex<3>(format, 'N')]; |
| 1146 | int c = in_dims[GetTensorDimIndex<3>(format, 'C')]; |
| 1147 | int d = in_dims[GetTensorDimIndex<3>(format, '0')]; |
| 1148 | int h = in_dims[GetTensorDimIndex<3>(format, '1')]; |
| 1149 | int w = in_dims[GetTensorDimIndex<3>(format, '2')]; |
| 1150 | |
| 1151 | // OneDNN requires dimensions in NCDHW format. |
| 1152 | return memory::dims({n, c, d, h, w}); |
| 1153 | } |
| 1154 | |
| 1155 | /// Map OneDNN memory::dims object into TensorShape object. |
| 1156 | /// |
no test coverage detected