| 198 | } |
| 199 | |
| 200 | int64 CalculateTensorSize(const OpInfo::TensorProperties& prop) { |
| 201 | int64 size = DataTypeSize(BaseType(prop.dtype())); |
| 202 | TensorShapeProto shape = prop.shape(); |
| 203 | |
| 204 | // Can't infer the size if the rank is unknown. It has to be at least a |
| 205 | // scalar though. |
| 206 | if (shape.unknown_rank()) { |
| 207 | VLOG(2) << "CalculateTensorSize() -- unknown rank"; |
| 208 | return size; |
| 209 | } |
| 210 | |
| 211 | // If one of the dimensions is unknown statically, assume it's at least one. |
| 212 | for (int i = 0; i < shape.dim_size(); ++i) { |
| 213 | if (shape.dim(i).size() < 0) { |
| 214 | shape.mutable_dim(i)->set_size(1); |
| 215 | VLOG(2) << "CalculateTensorSize() -- unknown dim: " << i; |
| 216 | } |
| 217 | } |
| 218 | |
| 219 | int64 num_elems = TensorShape(shape).num_elements(); |
| 220 | int64 tensor_size = MultiplyWithoutOverflow(num_elems, size); |
| 221 | if (tensor_size < 0) { |
| 222 | VLOG(1) << "Overflow encountered when computing tensor size, multiplying " |
| 223 | << num_elems << " with " << size; |
| 224 | return -1; |
| 225 | } |
| 226 | return tensor_size; |
| 227 | } |
| 228 | |
| 229 | int64 CalculateOutputSize( |
| 230 | const std::vector<OpInfo::TensorProperties>& output_properties, |