QEmbedding Class — pytorch Architecture

Source Code

aten/src/ATen/native/quantized/cpu/qembeddingbag.cpp lines 1491–1524

class QEmbedding final {
 public:
  static at::Tensor run(
      const c10::intrusive_ptr<EmbeddingPackedParamsBase>& packed_weight,
      const Tensor& indices,
      bool pruned_weights) {
    // Set default offsets here since the FBGEMM lookup op expects it.
    const auto offsets_size = indices.numel();
    at::Tensor offsets = at::arange(0, offsets_size, indices.scalar_type());
    at::Tensor output;
    static_assert(bit_rate==4 || bit_rate ==8,
          "Currently only support 8-bit embedding quantization");
    if (bit_rate == 8) {
      return packed_weight->embeddingbag_byte(
          indices,
          offsets,
          pruned_weights,
          std::nullopt,
          std::nullopt,
          false /* include_last_offset */,
          true /* is_embedding_op */);
    } else if (bit_rate == 4) {
      return packed_weight->embeddingbag_4bit(
          indices,
          offsets,
          pruned_weights,
          std::nullopt,
          std::nullopt,
          false,
          true);
    }
    return output;
  }
};