| 328 | // returned. |
| 329 | template <typename T, typename U> |
| 330 | tensorflow::Status NumElements(const std::vector<T>& shape, U* num_elements) { |
| 331 | static_assert( |
| 332 | std::numeric_limits<T>::max() <= std::numeric_limits<uint64_t>::max(), |
| 333 | "vector type exceed capabilities of NumElements"); |
| 334 | |
| 335 | *num_elements = 1; |
| 336 | for (const T& dim : shape) { |
| 337 | if (dim < 0) { |
| 338 | // TensorFlow's shapes sometimes include -1 to represent an "unknown" |
| 339 | // size but TOCO isn't able to create arrays of unknown sizes and will |
| 340 | // crash in RequiredBufferSizeForShape(). |
| 341 | return tensorflow::errors::InvalidArgument( |
| 342 | "Tensor shape should not include negative values"); |
| 343 | } |
| 344 | if (*num_elements != 0 && |
| 345 | static_cast<uint64_t>(dim) > |
| 346 | std::numeric_limits<U>::max() / *num_elements) { |
| 347 | *num_elements = 0; |
| 348 | return tensorflow::errors::InvalidArgument("Tensor shape is too large"); |
| 349 | } |
| 350 | *num_elements *= dim; |
| 351 | } |
| 352 | return tensorflow::Status::OK(); |
| 353 | } |
| 354 | |
| 355 | // A model file may have shuffled FC weights. |
| 356 | // When that happens, we want to de-shuffle them immediately on import, |