| 659 | template <typename Device, typename T, typename Index, |
| 660 | scatter_nd_op::UpdateOp Op> |
| 661 | Status DoScatterNd(OpKernelContext* c, const Tensor& indices, |
| 662 | const Tensor& updates, const TensorShape& shape, Tensor* out, |
| 663 | bool allocate) { |
| 664 | int64 slice_dim; |
| 665 | Index num_updates; |
| 666 | Index slice_size; |
| 667 | TF_RETURN_IF_ERROR(PrepareAndValidateInputs<Index>( |
| 668 | shape, indices, updates, &slice_dim, &num_updates, &slice_size)); |
| 669 | |
| 670 | IndexFlattener<Device, Index> index_flattener; |
| 671 | auto indices_flat = index_flattener(c, indices); |
| 672 | auto updates_flat = updates.shaped<T, 2>({num_updates, slice_size}); |
| 673 | |
| 674 | if (allocate) { |
| 675 | TF_RETURN_IF_ERROR(c->allocate_temp(DataTypeToEnum<T>::value, shape, out)); |
| 676 | } else { |
| 677 | CHECK_NOTNULL(out); |
| 678 | } |
| 679 | |
| 680 | if (shape.num_elements() == 0) { |
| 681 | return Status::OK(); |
| 682 | } |
| 683 | |
| 684 | if (allocate) { |
| 685 | // Brand new tensor, zero it out. |
| 686 | functor::SetZeroFunctor<Device, T> fill; |
| 687 | fill(c->eigen_device<Device>(), out->flat<T>()); |
| 688 | } |
| 689 | auto output_matrix = |
| 690 | out->shaped<T, 2>({shape.num_elements() / slice_size, slice_size}); |
| 691 | |
| 692 | Index bad_i = -1; |
| 693 | |
| 694 | if (shape.num_elements() > 0) { |
| 695 | switch (slice_dim) { |
| 696 | #define PARAMS_CASE(IXDIM) \ |
| 697 | case IXDIM: { \ |
| 698 | typename Eigen::array<Eigen::DenseIndex, IXDIM> output_shape_prefix; \ |
| 699 | for (int i = 0; i < IXDIM; ++i) { \ |
| 700 | output_shape_prefix[i] = shape.dim_size(i); \ |
| 701 | } \ |
| 702 | functor::ScatterNdFunctor<Device, T, Index, Op, IXDIM> functor; \ |
| 703 | bad_i = \ |
| 704 | functor(c->eigen_device<Device>(), slice_size, output_shape_prefix, \ |
| 705 | output_matrix, indices_flat, updates_flat, output_matrix); \ |
| 706 | } break |
| 707 | // TODO(simister): Re-enable this once binary size is under control. |
| 708 | // PARAMS_CASE(0); |
| 709 | PARAMS_CASE(1); |
| 710 | PARAMS_CASE(2); |
| 711 | PARAMS_CASE(3); |
| 712 | PARAMS_CASE(4); |
| 713 | PARAMS_CASE(5); |
| 714 | PARAMS_CASE(6); |
| 715 | PARAMS_CASE(7); |
| 716 | #undef PARAMS_CASE |
| 717 | default: |
| 718 | return errors::InvalidArgument( |
nothing calls this directly
no test coverage detected