Source code
Revision control
Copy as Markdown
Other Tools
#include "pos_enc.hpp"
#include <cmath>
#include <cassert>
namespace pk {
// INF_VAL from NeMo multi_head_attention.py (used as the sinusoid base).
static constexpr double kInfVal = 10000.0;
void rel_pos_encoding(int T, int d_model, std::vector<float>& out) {
assert(T > 0 && d_model > 0 && (d_model % 2) == 0);
const int P = 2 * T - 1; // number of relative positions
const int half = d_model / 2; // number of (sin,cos) pairs
// div_term[i] = exp(2i * -(log(INF_VAL)/d_model)) for i in [0, half).
// Compute in double for parity, store the precomputed factors.
std::vector<double> div_term(half);
const double factor = -(std::log(kInfVal) / (double)d_model);
for (int i = 0; i < half; ++i) {
div_term[i] = std::exp((double)(2 * i) * factor);
}
out.assign((size_t)P * d_model, 0.0f);
// positions run from +(T-1) down to -(T-1).
for (int p = 0; p < P; ++p) {
const double pos = (double)((T - 1) - p); // (T-1), (T-2), ..., -(T-1)
float* row = out.data() + (size_t)p * d_model;
for (int i = 0; i < half; ++i) {
const double arg = pos * div_term[i];
row[2 * i] = (float)std::sin(arg); // even dims
row[2 * i + 1] = (float)std::cos(arg); // odd dims
}
}
}
void local_rel_pos_encoding(int att_left, int att_right, int d_model,
std::vector<float>& out) {
assert(att_left >= 0 && att_right >= 0 && d_model > 0 && (d_model % 2) == 0);
const int P = att_left + att_right + 1; // local relative positions
const int half = d_model / 2;
std::vector<double> div_term(half);
const double factor = -(std::log(kInfVal) / (double)d_model);
for (int i = 0; i < half; ++i) {
div_term[i] = std::exp((double)(2 * i) * factor);
}
out.assign((size_t)P * d_model, 0.0f);
// positions run from +att_left down to -att_right (NeMo arange(left,-right-1,-1)).
for (int p = 0; p < P; ++p) {
const double pos = (double)(att_left - p);
float* row = out.data() + (size_t)p * d_model;
for (int i = 0; i < half; ++i) {
const double arg = pos * div_term[i];
row[2 * i] = (float)std::sin(arg);
row[2 * i + 1] = (float)std::cos(arg);
}
}
}
} // namespace pk