32 const Shape &coords_shape,
const CoordsT *coords_data,
const Shape &,
35 int axis = op_params.
axis;
42 const int axis_size = input_shape.
Dims(axis);
43 const int coords_count = coords_shape.
FlatSize();
46 for (
int i = 0; i < axis; ++i)
48 outer_size *= input_shape.
Dims(i);
54 inner_size *= input_shape.
Dims(i);
57 for (
int outer = 0; outer < outer_size; ++outer)
59 for (
int i = 0; i < coords_count; ++i)
61 assert(coords_data[i] >= 0);
62 assert(coords_data[i] < axis_size);
63 std::memcpy(output_data + (outer * coords_count + i) * inner_size,
64 input_data + (outer * axis_size + coords_data[i]) * inner_size,
65 sizeof(T) * inner_size);
void Gather(const GatherParams &op_params, const Shape &input_shape, const T *input_data, const Shape &coords_shape, const CoordsT *coords_data, const Shape &, T *output_data)