Function GetTensorDimsFromSpatialDims

tensorflow/core/util/tensor_format.h:143–158

Returns the rank of a tensor with 'num_spatial_dims' spatial dimensions and tensor format 'format'. This is the inverse of GetTensorSpatialDims.

// Returns the rank of a tensor with 'num_spatial_dims' spatial dimensions and
// tensor format 'format'. This is the inverse of GetTensorSpatialDims.
inline int GetTensorDimsFromSpatialDims(int num_spatial_dims,
                                        TensorFormat format) {
  switch (format) {
    case FORMAT_NHWC:
    case FORMAT_NCHW:
    case FORMAT_HWNC:
    case FORMAT_HWCN:
      return num_spatial_dims + 2;  // Include N,C.
    case FORMAT_NCHW_VECT_C:
    case FORMAT_NHWC_VECT_W:
      return num_spatial_dims + 3;  // Include N,C,VectDim.
    default:
      LOG(FATAL) << "Unknown format " << format;
      return -1;  // Avoid compiler warning about missing return value
  }
}

// Returns the rank of a tensor with 'num_spatial_dims' spatial dimensions and
// filter tensor format 'format'.