| #pragma once | |
| #include <cstdint> | |
| namespace caffe2 { | |
| // clang-format off | |
| /** | |
| * Embedding lookup with reduction. | |
| * | |
| * `input` of size data_size * block_size | |
| * `indices` of size index_size | |
| * `offsets` of size output_size | |
| * `weights` nullptr or array of size index_size | |
| * `out` of size output_size * block_size | |
| * | |
| * Behavior is roughly equivalent to pseudocode: | |
| * | |
| * pos = 0 | |
| * for (i = 0..output_size-1) | |
| * for (k = 0..block_size-1) | |
| * out[i*block_size + k] = 0 | |
| * start_offset = offsets[i] | |
| * end_offset = offsets[i+1] | |
| * length = end_offset - start_offset | |
| * for (j = start_offset..end_offset-1) | |
| * for (k = 0..block_size-1) | |
| * out[i*block_size + k] += input[indices[pos]*block_size + k] * | |
| * (weights ? weights[IS_WEIGHT_POSITIONAL ? j - start_offset : pos] : 1.0) | |
| * pos += 1 | |
| * if (normalize_weights && length > 0) | |
| * for (k = 0..block_size-1) | |
| * out[i*block_size + k] /= length | |
| * | |
| * TODO: make this API also take "offsets" rather than "lengths" to match the | |
| * API for PyTorch's EmbeddingBag | |
| */ | |
| // clang-format on | |
| template < | |
| typename IndexType, | |
| typename InType, | |
| typename OutType, | |
| bool IS_WEIGHT_POSITIONAL = false> | |
| void EmbeddingLookupIdx( | |
| const std::int64_t block_size, | |
| const std::int64_t output_size, | |
| const std::int64_t index_size, | |
| const std::int64_t data_size, | |
| const InType* input, | |
| const IndexType* indices, | |
| const IndexType* offsets, | |
| const float* weights, // optional, can be null for non-weighted sum | |
| const float* scale_bias, // optional scale & bias params for uint8 input | |
| bool normalize_by_lengths, | |
| OutType* out); | |
| } // namespace caffe2 |