| 190 | |
| 191 | template <class Shape> |
| 192 | Status TensorShapeBase<Shape>::InitDims(gtl::ArraySlice<int64> dim_sizes) { |
| 193 | DCHECK_EQ(tag(), REP16); |
| 194 | |
| 195 | // Allow sizes that are under kint64max^0.25 so that 4-way multiplication |
| 196 | // below cannot overflow. |
| 197 | static const uint64 kMaxSmall = 0xd744; |
| 198 | static_assert(kMaxSmall * kMaxSmall * kMaxSmall * kMaxSmall <= kint64max, |
| 199 | "bad overflow check"); |
| 200 | bool large_size = false; |
| 201 | for (auto s : dim_sizes) { |
| 202 | if (s > kMaxSmall) { |
| 203 | large_size = true; |
| 204 | break; |
| 205 | } |
| 206 | } |
| 207 | |
| 208 | if (!kIsPartial && !large_size) { |
| 209 | for (auto s : dim_sizes) { |
| 210 | if (TF_PREDICT_FALSE(s < 0)) { |
| 211 | return errors::InvalidArgument( |
| 212 | "Expected shape dimensions to be non-negative, got ", s); |
| 213 | } |
| 214 | } |
| 215 | } |
| 216 | |
| 217 | if (!large_size) { |
| 218 | // Every size fits in 16 bits; use fast-paths for dims in {1,2,3,4}. |
| 219 | uint16* dst = as16()->dims_; |
| 220 | switch (dim_sizes.size()) { |
| 221 | case 1: { |
| 222 | set_ndims_byte(1); |
| 223 | const int64 size = dim_sizes[0]; |
| 224 | const bool neg = Set16(kIsPartial, dst, 0, size); |
| 225 | set_num_elements(neg ? -1 : size); |
| 226 | return Status::OK(); |
| 227 | } |
| 228 | case 2: { |
| 229 | set_ndims_byte(2); |
| 230 | const int64 size0 = dim_sizes[0]; |
| 231 | const int64 size1 = dim_sizes[1]; |
| 232 | bool neg = Set16(kIsPartial, dst, 0, size0); |
| 233 | neg |= Set16(kIsPartial, dst, 1, size1); |
| 234 | set_num_elements(neg ? -1 : (size0 * size1)); |
| 235 | return Status::OK(); |
| 236 | } |
| 237 | case 3: { |
| 238 | set_ndims_byte(3); |
| 239 | const int64 size0 = dim_sizes[0]; |
| 240 | const int64 size1 = dim_sizes[1]; |
| 241 | const int64 size2 = dim_sizes[2]; |
| 242 | bool neg = Set16(kIsPartial, dst, 0, size0); |
| 243 | neg |= Set16(kIsPartial, dst, 1, size1); |
| 244 | neg |= Set16(kIsPartial, dst, 2, size2); |
| 245 | set_num_elements(neg ? -1 : (size0 * size1 * size2)); |
| 246 | return Status::OK(); |
| 247 | } |
| 248 | case 4: { |
| 249 | set_ndims_byte(4); |
no test coverage detected