[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.

[Copilot]

Pull request overview

Copilot reviewed 5 out of 5 changed files in this pull request and generated no new comments.

Suppressed comments (14)

src/embkit/encoding/init.py:122

• self.shape is inconsistent with the runtime dim calculation when full_len is None: it currently adds both +1 and +pe_dim for positional encoding, but dim only adds pe_dim. This makes the advertised output shape incorrect.

        if self.full_len is not None:
            self.shape = (self.full_len, len(self.alphabet) + (self.pe_dim if self.encode_pos else 0)) 
        else:
            self.shape = (len(self.alphabet) + (1 if self.encode_pos else 0) + self.pe_dim,)

src/embkit/encoding/init.py:155

• The per-call dim calculation should match the self.shape logic (including the special-case for a single scalar position channel when pe_dim == 1). Right now dim always adds pe_dim, which diverges from the intended scalar position mode and can desync with self.shape.

        dim = len(self.alphabet) + (self.pe_dim if self.encode_pos else 0)
        batch_size = len(seqs)
        if self._use_numpy_backend():
            np_dtype = self.np_dtype or np.float32
            one_hot_matrix = np.zeros((batch_size, FL, dim), dtype=np_dtype)
        else:
            one_hot_matrix = torch.zeros((batch_size, FL, dim), device=self.device, dtype=self.torch_dtype)

src/embkit/encoding/init.py:213

• to_dict() now includes pe_dim, but from_dict() ignores it. This breaks round-tripping and makes it impossible to restore a non-default positional encoding dimensionality.

    def from_dict(cls, data):
        dtype_str = data.get("dtype", "float32")
        if dtype_str == "float32":
            dtype = np.float32
        elif dtype_str == "float64":
            dtype = np.float64
        else:
            dtype = None  # default to torch's default dtype
        return cls(
            full_len=data.get("full_len"),
            encode_x=data.get("encode_x", True),
            encode_pos=data.get("encode_pos", False),
            device=data.get("device"),
            dtype=dtype,
            backend=data.get("backend", 'torch')
        )

src/embkit/encoding/init.py:285

• PreEncoded.__call__ silently returns None for unsupported backends. It should raise, so failures are obvious and don't propagate as confusing downstream errors.

    def __call__(self, names):
        if isinstance(names, str):
            return self.cache[names]
        if self.backend == "numpy":
            return np.array([self.cache[n] for n in names], dtype=np.float32)
        elif self.backend == "torch":
            return torch.stack([self.cache[n] for n in names])

src/embkit/modules/mha.py:124

• Same as above: from_dict should typically instantiate via cls(...) rather than MHAPooling(...) to preserve subclassing/renames and follow project conventions.

    def from_dict(cls, params):
        return MHAPooling(
            embed_dim=params["embed_dim"],
            num_heads=params["num_heads"]
        )

src/embkit/encoding/init.py:85

• Defaulting pe_dim to 2 changes the encode_pos=True output dimensionality (and breaks existing expectations/tests that add a single position column). Consider defaulting to pe_dim=1 for backward compatibility, and treat pe_dim>1 as the sinusoidal mode.

    def __init__(self, full_len=None, encode_x=True, encode_pos=False, pe_dim=2, device=None, dtype=torch.float32, backend='torch'):

src/embkit/encoding/init.py:177

• When encode_pos is enabled, the scalar position value written at len(self.alphabet) is immediately overwritten by the sinusoidal fill (one_hot_matrix[b, :, -self.pe_dim:]). This makes the scalar mode ineffective and causes surprising outputs for pe_dim==1.

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

                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)
            if self.encode_pos and self.pe_dim > 0 and not self._use_numpy_backend():
                # Fill PE channels for this sequence
                pe = torch.stack([position_sin_cos_tensor(i, self.pe_dim, device=self.device, dtype=self.torch_dtype) for i in range(FL)], dim=0)
                one_hot_matrix[b, :, -self.pe_dim:] = pe
            elif self.encode_pos and self.pe_dim > 0 and self._use_numpy_backend():
                # numpy fallback: compute via torch then convert
                pe = np.stack([position_sin_cos(i, self.pe_dim) for i in range(FL)], axis=0)
                one_hot_matrix[b, :, -self.pe_dim:] = pe

src/embkit/encoding/init.py:266

• PreEncoded accepts any backend value, but unsupported values lead to an empty cache and shape with dim=None. Validate backend early and fail fast with a clear error.

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

    def __init__(self, path, backend="numpy"):
        self.path = path
        self.backend = backend
        reader = CsvReader(path, index_column=0, header=None, sep="\t")
        self.cache = {}

src/embkit/datasets/init.py:69

• DataFrameMapper assumes each mapper returns a torch Tensor (calls .to(...)). If a mapper returns a numpy array / list / scalar, this will raise AttributeError. Converting non-tensors via torch.as_tensor makes the dataset wrapper more robust.

    def __getitem__(self, idx):
        row = self.data.iloc[idx]
        out = []
        for k, v in self.mappers:
            out.append( v(row[k]).to(self.device, dtype=self.dtype) )
        return out

src/embkit/optimize/multitask.py:128

• When device is provided, batches are moved to that device, but the task models are not. If callers forget to move models themselves, this will trigger device-mismatch errors (e.g., CUDA inputs with CPU model). Consider moving all task models onto device inside the training helper.

    loaders, trainable_params = _prepare_learning_tasks(tasks)

    optimizer = Adam(trainable_params, lr=lr)
    scheduler = StepLR(optimizer, step_size=1, gamma=0.5)

src/embkit/optimize/multitask.py:196

• Same device-mismatch risk as in multi_task_train_weighted_sync: inputs are moved to device, but models are not. Moving task models once up-front avoids runtime errors and makes the device argument behave as expected.

    loaders, trainable_params = _prepare_learning_tasks(tasks)
    normalized_schedule = _normalize_task_schedule(task_schedule, len(tasks))

    optimizer = Adam(trainable_params, lr=lr)
    scheduler = StepLR(optimizer, step_size=1, gamma=0.5)

src/embkit/models/pair.py:25

• device is applied only to self.pair_predict, leaving item_module and context_module on their original devices. This can cause device-mismatch errors at runtime when concatenating outputs. Prefer moving the whole module (including submodules) via self.to(...).

        self.pair_predict = nn.Sequential(
            nn.BatchNorm1d( input_dim, dtype=dtype),
            nn.Linear( input_dim, learning_dim, dtype=dtype ), nn.ReLU(),
            nn.Dropout(0.2),
            nn.Linear( learning_dim, learning_dim, dtype=dtype), nn.ReLU(),
            # nn.Linear( learning_dim, learning_dim, dtype=dtype), nn.ReLU(),
            nn.Linear( learning_dim, 1, dtype=dtype)
        ).to(device)

src/embkit/models/pair.py:46

• from_dict is a @classmethod but returns a concrete PairPredictor(...) rather than cls(...). Using cls matches the convention used elsewhere (e.g. src/embkit/models/ffnn.py:67) and preserves subclassing/renames.

    def from_dict(cls, params):
        return PairPredictor(
            item_module=factory.build(params["item_module"]),
            context_module=factory.build(params["context_module"]),
            item_dim=params["item_dim"],
            context_dim=params["context_dim"],
            learning_dim=params["learning_dim"]
        )

src/embkit/modules/mha.py:60

• from_dict is a @classmethod but returns a concrete MHABlock(...) rather than cls(...). Using cls matches the convention used elsewhere (e.g. src/embkit/models/ffnn.py:67) and avoids issues if the class is renamed or subclassed.

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

    def from_dict(cls, params):
        return MHABlock(
            embed_dim=params["embed_dim"],
            num_heads=params["num_heads"],
            dropout=params["dropout"],
            out_dim=params["out_dim"],
            ffn_inner_dim=params["ffn_inner_dim"]
        )