| 551 | |
| 552 | template <typename Index> |
| 553 | Status PrepareAndValidateInputs(const TensorShape& params_shape, |
| 554 | const Tensor& indices, const Tensor& updates, |
| 555 | int64* slice_dim, Index* num_updates, |
| 556 | Index* slice_size) { |
| 557 | const TensorShape& indices_shape(indices.shape()); |
| 558 | const TensorShape& updates_shape(updates.shape()); |
| 559 | |
| 560 | if (!TensorShapeUtils::IsVectorOrHigher(params_shape)) { |
| 561 | return errors::InvalidArgument("Output must be at least 1-D, ", |
| 562 | "got shape: ", params_shape.DebugString()); |
| 563 | } |
| 564 | |
| 565 | if (!ValidEmptyOutputShape(params_shape.num_elements(), |
| 566 | indices_shape.num_elements(), |
| 567 | updates_shape.num_elements())) { |
| 568 | return errors::InvalidArgument( |
| 569 | "Indices and updates specified for empty output. indices shape: ", |
| 570 | indices.shape().DebugString()); |
| 571 | } |
| 572 | |
| 573 | if (updates.dim_size(0) != indices.dim_size(0)) { |
| 574 | return errors::InvalidArgument( |
| 575 | "The outermost dimension of updates and indices ", |
| 576 | "must match. Got indices.shape ", indices_shape.DebugString(), |
| 577 | ", updates.shape ", updates_shape.DebugString()); |
| 578 | } |
| 579 | TF_RETURN_IF_ERROR(ValidateUpdateShape(params_shape, indices, updates)); |
| 580 | |
| 581 | // Check that we have enough index space |
| 582 | const int64 N_big = indices.NumElements(); |
| 583 | if (N_big > std::numeric_limits<Index>::max()) { |
| 584 | return errors::InvalidArgument("indices has too many elements for ", |
| 585 | DataTypeString(DataTypeToEnum<Index>::v()), |
| 586 | " indexing: ", N_big, " > ", |
| 587 | std::numeric_limits<Index>::max()); |
| 588 | } |
| 589 | if (params_shape.dim_size(0) > std::numeric_limits<Index>::max()) { |
| 590 | return errors::InvalidArgument("params_shape[0] too large for ", |
| 591 | DataTypeString(DataTypeToEnum<Index>::v()), |
| 592 | " indexing: ", params_shape.dim_size(0), |
| 593 | " > ", std::numeric_limits<Index>::max()); |
| 594 | } |
| 595 | |
| 596 | // Calculate the number of dimensions in indices |
| 597 | *slice_dim = (indices_shape.dims() > 1) |
| 598 | ? indices_shape.dim_size(indices_shape.dims() - 1) |
| 599 | : 1; |
| 600 | |
| 601 | // Calculate the number of elements that make up each slice of our updated |
| 602 | // tensor. This allows us to work with flattened tensors and copy over whole |
| 603 | // slices at a time. |
| 604 | Index total_nd = params_shape.dims(); |
| 605 | |
| 606 | int64 slice_size_big = 1; |
| 607 | for (int64 i = *slice_dim; i < total_nd; ++i) { |
| 608 | slice_size_big *= params_shape.dim_size(i); |
| 609 | } |
| 610 |
nothing calls this directly
no test coverage detected