Check whether updates.shape = indices.shape[:batch_dim] + params_shape[slice_dim:]
| 515 | // Check whether updates.shape = indices.shape[:batch_dim] + |
| 516 | // params_shape[slice_dim:] |
| 517 | Status ValidateUpdateShape(const TensorShape& params_shape, |
| 518 | const Tensor& indices, const Tensor& updates) { |
| 519 | const int64 slice_dim = |
| 520 | (indices.dims() > 1) ? indices.dim_size(indices.dims() - 1) : 1; |
| 521 | const int64 batch_dim = (indices.dims() > 1) ? indices.dims() - 1 : 1; |
| 522 | |
| 523 | auto shape_err = [&]() { |
| 524 | return errors::InvalidArgument( |
| 525 | "Must have updates.shape = indices.shape[:batch_dim] + ", |
| 526 | "params_shape[slice_dim:], got updates.shape: ", |
| 527 | updates.shape().DebugString(), |
| 528 | ", indices.shape: ", indices.shape().DebugString(), |
| 529 | ", params_shape: ", params_shape.DebugString(), |
| 530 | ", slice_dim: ", slice_dim, ", and batch_dim: ", batch_dim); |
| 531 | }; |
| 532 | |
| 533 | if (updates.dims() < batch_dim) return shape_err(); |
| 534 | if (params_shape.dims() < slice_dim + (updates.dims() - batch_dim)) { |
| 535 | return shape_err(); |
| 536 | } |
| 537 | if (updates.dims() != batch_dim + params_shape.dims() - slice_dim) { |
| 538 | return shape_err(); |
| 539 | } |
| 540 | for (int d = 0; d < batch_dim; ++d) { |
| 541 | if (updates.dim_size(d) != indices.dim_size(d)) return shape_err(); |
| 542 | } |
| 543 | for (int d = 0; d < updates.dims() - batch_dim; ++d) { |
| 544 | if (updates.dim_size(d + batch_dim) != |
| 545 | params_shape.dim_size(d + slice_dim)) { |
| 546 | return shape_err(); |
| 547 | } |
| 548 | } |
| 549 | return Status::OK(); |
| 550 | } |
| 551 | |
| 552 | template <typename Index> |
| 553 | Status PrepareAndValidateInputs(const TensorShape& params_shape, |
no test coverage detected