[WIP] More Encoding

[kellrott]

• Positional encoding
• Unit testing for positional encoding dimensions

[Copilot]

Pull request overview

This WIP PR extends the embkit.encoding module to support richer positional encodings (sin/cos) for ProteinOneHotEncoder, and adds a dtype option to OneHotEncoder’s precomputed one-hot tensors.

Changes:

• Added dtype parameter support to OneHotEncoder initialization and mapping construction.
• Added pe_dim to ProteinOneHotEncoder and introduced sinusoidal positional encoding helpers (position_sin_cos*).
• Updated ProteinOneHotEncoder (de)serialization to include the new positional-encoding configuration and added PE helper functions.

Suppressed comments (3)

src/embkit/encoding/init.py:16

• OneHotEncoder accepts dtype, but it isn’t stored on the instance, making it hard to apply consistently (e.g., for batch outputs).

This issue also appears on line 19 of the same file.

    def __init__(self, classes, device=None, dtype=None):
        self.classes = sorted(classes)
        self.num_classes = len(self.classes)
        self.mapping = {}
        self.class_idx = {}

src/embkit/encoding/init.py:168

• Inside the residue loop, a scalar position value is written to one_hot_matrix[..., len(self.alphabet)], but for pe_dim>0 those channels are overwritten immediately by the sinusoidal PE assignment below. Removing this avoids redundant work and prevents accidental out-of-bounds writes if pe_dim is changed.

                if self.encode_pos:
                    if self.full_len is not None:
                        one_hot_matrix[b, i, len(self.alphabet)] = float(i) / float(self.full_len)
                    else:
                        one_hot_matrix[b, i, len(self.alphabet)] = float(i)

src/embkit/encoding/init.py:21

• dtype is now applied to the precomputed per-class tensors (self.mapping[...]), but the batch path in __call__ still returns the default one_hot dtype (typically int64). This makes single-label and batch outputs inconsistent when dtype is provided.

        for i, n in enumerate(self.classes):
            self.mapping[n] = F.one_hot( torch.tensor(i), self.num_classes ).to(device=device, dtype=dtype)
            self.class_idx[n] = i

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.