Function gather

tests/validation/reference/Gather.cpp:41–99

{
template <typename T>
SimpleTensor<T> gather(const SimpleTensor<T> &src, const SimpleTensor<uint32_t> &indices, uint32_t actual_axis)
{
    const TensorShape dst_shape =
        arm_compute::misc::shape_calculator::compute_gather_shape(src.shape(), indices.shape(), actual_axis);
    SimpleTensor<T> dst(dst_shape, src.data_type());

    const auto src_ptr     = static_cast<const T *>(src.data());
    const auto indices_ptr = static_cast<const uint32_t *>(indices.data());
    const auto dst_ptr     = static_cast<T *>(dst.data());

    const uint32_t index_limit = src.shape()[actual_axis];

    Window win;
    win.use_tensor_dimensions(dst_shape);

    execute_window_loop(win,
                        [&](const Coordinates &dst_coords)
                        {
                            const auto dst_addr = coords2index(dst.shape(), dst_coords);

                            // Calculate the coordinates of the index value.
                            Coordinates idx_coords;

                            for (size_t i = 0; i < indices.shape().num_dimensions(); ++i)
                            {
                                idx_coords.set(i, dst_coords[i + actual_axis]);
                            }

                            const auto index = indices_ptr[coords2index(indices.shape(), idx_coords)];

                            if (index < index_limit)
                            {
                                // Calculate the coordinates of the source data.
                                Coordinates src_coords;

                                for (size_t i = 0; i < actual_axis; ++i)
                                {
                                    src_coords.set(i, dst_coords[i]);
                                }

                                src_coords.set(actual_axis, index);

                                for (size_t i = actual_axis + 1; i < src.shape().num_dimensions(); ++i)
                                {
                                    src_coords.set(i, dst_coords[i + indices.shape().num_dimensions() - 1]);
                                }

                                // Copy the data.
                                const auto src_addr = coords2index(src.shape(), src_coords);
                                dst_ptr[dst_addr]   = src_ptr[src_addr];
                            }
                            else
                            {
                                dst_ptr[dst_addr] = 0;
                            }
                        });

    return dst;