| 311 | } |
| 312 | |
| 313 | DimensionHandle InferenceContext::NumElements(ShapeHandle s) { |
| 314 | const auto rank = Rank(s); |
| 315 | if (rank == kUnknownRank) return UnknownDim(); |
| 316 | bool found_unknown = false; |
| 317 | int64 size = 1; |
| 318 | for (int i = 0; i < rank; ++i) { |
| 319 | int64 dim_val = Value(Dim(s, i)); |
| 320 | if (dim_val == kUnknownDim) { |
| 321 | found_unknown = true; |
| 322 | } else if (dim_val == 0) { |
| 323 | return MakeDim(0); |
| 324 | } else { |
| 325 | size *= dim_val; |
| 326 | } |
| 327 | } |
| 328 | if (found_unknown) { |
| 329 | return UnknownDim(); |
| 330 | } else { |
| 331 | return MakeDim(size); |
| 332 | } |
| 333 | } |
| 334 | |
| 335 | string InferenceContext::DebugString(ShapeHandle s) { |
| 336 | if (RankKnown(s)) { |