pos_enc.hpp - mozsearch

Enable keyboard shortcuts

#pragma once

#include <vector>

namespace pk {

// Relative positional encoding table for the FastConformer encoder, matching

// NeMo's RelPositionalEncoding (multi_head_attention.py).

//

// For an input of length T, NeMo builds `pos_emb` over 2T-1 relative positions

// running from +(T-1) DOWN TO -(T-1) (inclusive). Each row p holds a sinusoid:

//   div_term[i] = exp(2i * -(log(10000)/d_model))           for i in [0, d_model/2)

//   pe[p, 2i]   = sin(position(p) * div_term[i])

//   pe[p, 2i+1] = cos(position(p) * div_term[i])

// where position(0) = T-1, position(1) = T-2, ..., position(2T-2) = -(T-1).

//

// (NeMo precomputes the table over a larger length and slices the center; the

// sliced positions are exactly +(T-1)..-(T-1), independent of the precompute

// length, so computing directly for T reproduces the same values.)

//

// Output: row-major [2T-1, d_model] (d_model fastest), i.e. out[p*d_model + c].

void rel_pos_encoding(int T, int d_model, std::vector<float>& out);

// LOCAL relative positional encoding for NeMo rel_pos_local_attn

// (LocalAttRelPositionalEncoding): positions run from +att_left DOWN TO

// -att_right (att_left+att_right+1 rows), using the SAME sinusoid as

// rel_pos_encoding. These are exactly the centre rows of the full table, so

// banded attention's positional term is bit-identical to NeMo's local pos.

//

// Output: row-major [att_left+att_right+1, d_model] (d_model fastest).

void local_rel_pos_encoding(int att_left, int att_right, int d_model,

                            std::vector<float>& out);

} // namespace pk