Foundations

Categorical and numeric molecular feature encoding

MediumRaw PyTorchAbout 20 min1 public example plus private tests

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Your task

Encode categorical fields with explicit unknown-value handling and combine them with transformed numeric fields. Return categorical indicators, combined features, and the supplied projection.

Function signature

categorical_feature_encoding(
    categories: torch.Tensor,
    choices: torch.Tensor,
    continuous: torch.Tensor,
    continuous_transform: torch.Tensor,
    continuous_scale: torch.Tensor,
    parameters: dict,
    masks: dict,
    unknown_policy: str = 'bucket',
) -> dict[str, torch.Tensor]

Worked example

Known, unknown, and numerical features

For choices [6, 7, 8], category 6 uses the first channel, category 7 the second, and category 9 the extra unknown channel. In the second field, -1 is an explicitly listed category, so it is known. A missing field would instead have all-zero channels. Numerical values [12, 20, 14] with scale $0.01$ become [0.12, 0.20, 0.14]; the identity projection preserves the assembled features.

import torch

categories = torch.tensor([[[6, -1], [9, 0], [7, 4]]], dtype=torch.int64)
choices = torch.tensor([[6, 7, 8], [-1, 0, 1]], dtype=torch.int64)
continuous = torch.tensor([[[12.0, 0.0], [20.0, 1.0], [14.0, 0.0]]], dtype=torch.float64, requires_grad=True)
continuous_transform = torch.tensor([0, 0], dtype=torch.int64)
continuous_scale = torch.tensor([0.01, 1.0], dtype=torch.float64, requires_grad=True)
parameters = {
    'weight': torch.tensor([[1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.0], [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0]], dtype=torch.float64, requires_grad=True),
    'bias': torch.full((10,), 0.0, dtype=torch.float64, requires_grad=True),
}
masks = {
    'choice_mask': torch.tensor([[True, True, True], [True, True, True]], dtype=torch.bool),
    'field_mask': torch.tensor([[[True, True], [True, True], [True, True]]], dtype=torch.bool),
    'mask': torch.tensor([[True, True, True]], dtype=torch.bool),
}
unknown_policy = 'bucket'

result = categorical_feature_encoding(categories=categories, choices=choices, continuous=continuous, continuous_transform=continuous_transform, continuous_scale=continuous_scale, parameters=parameters, masks=masks, unknown_policy=unknown_policy)

Expected return

{
    'one_hot': torch.tensor([[[[1.0, 0.0, 0.0, 0.0], [1.0, 0.0, 0.0, 0.0]], [[0.0, 0.0, 0.0, 1.0], [0.0, 1.0, 0.0, 0.0]], [[0.0, 1.0, 0.0, 0.0], [0.0, 0.0, 0.0, 1.0]]]], dtype=torch.float64),
    'features': torch.tensor([[[1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.12, 0.0], [0.0, 0.0, 0.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.2, 1.0], [0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.14, 0.0]]], dtype=torch.float64),
    'embedding': torch.tensor([[[1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 0.12, 0.0], [0.0, 0.0, 0.0, 1.0, 0.0, 1.0, 0.0, 0.0, 0.2, 1.0], [0.0, 1.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0, 0.14, 0.0]]], dtype=torch.float64),
}

Parameters

Parameter	Shape and meaning
`categories`torch.Tensor	$[\ldots, F]$ Prepared raw category value for each field.
`choices`torch.Tensor	$[F, K]$ Allowed category values in each field’s vocabulary order.
`continuous`torch.Tensor	$[\ldots, D]$ `float` $[\ldots,D]$ , prepared raw numeric descriptors, $D\ge 0$ .
`continuous_transform`torch.Tensor	$[D]$ `int64` $[D]$ , code $0=\mathrm{identity}$ or $1=\operatorname{asinh}$ per numeric field.
`continuous_scale`torch.Tensor	$[D]$ `float` $[D]$ , per-field multiplier after identity/asinh.
`parameters`dict	Supplied model parameters, keyed as shown below.
`masks`dict	Supplied masks. Each field below defines what True includes or permits.
`unknown_policy`strDefault `'bucket'`	"bucket" assigns an unmatched valid category to its field’s extra unknown bucket; "zero" leaves that field’s categorical block all zero.

Return

dict[str, torch.Tensor]

Return exactly the named fields applicable to this operation.

Required behavior

Map multiple integer descriptor fields to explicit one-hots, concatenate numeric descriptors, and project the assembled feature. categories $[\ldots,F]$ is int64. choices $[F,K]$ lists allowed raw values (including legitimate negative values such as formal charge); choice_mask $[F,K]$ marks real vocabulary slots.
Valid choices within each field must be unique; unused slots remain zero feature columns. $F,K\ge 1$ . field_mask $[\ldots,F]$ marks present fields; mask $[\ldots]$ marks valid observations. A missing field has ALL zero category channels, distinct from an unknown but present value.
For each valid/present field, compare its raw value with all valid choices. Matched slot $k$ sets category channel $k$ to $1$ . If no match, unknown_policy='bucket' (default) sets the last channel $K$ to $1$ ; 'zero' leaves every channel zero. Return one_hot $[\ldots,F,K+1]$ in continuous.dtype.
Each numeric descriptor continuous $[\ldots,D]$ ( $D$ may be $0$ ) has a supplied continuous_transform $[D]$ : code $0$ keeps $x$ ; code $1$ uses $\operatorname{asinh}(x)$ . Multiply this transformed value by continuous_scale $[D]$ .
This covers Chemprop's mass factor $0.01$ and AF3's asinh reference charge (ESMFold2 uses raw charge), without importing a ready-made asinh. Flatten one_hot's last two axes in field-major order and concatenate the transformed/scaled numeric descriptors, calling this features $[\ldots,F\cdot (K+1)+D]$ .
Then $\mathrm{embedding}=\mathrm{features}\mathbin{@}\mathrm{weight}+\mathrm{bias}$ , with weight $[F\cdot (K+1)+D,O]$ ,bias $[O]$ , $O\ge 1$ . Masked rows have zero features and zero embedding EVEN with nonzero bias. Masked continuous payloads may be NaN and must be sanitized before arithmetic.
Unknown unknown_policy values raise ValueError. Floating tensors use one float32/float64 dtype. Observed descriptors and parameters are finite, continuous_transform codes are 0 or 1, and numeric magnitudes are below 10^10. Scaling and projection intermediates must remain finite in that dtype.
Preserve the smooth derivative $1/\sqrt{1+x^2}$ , including derivative 1 at x=0. The supplied fields and vocabularies are ready to use; no chemical parsing, bond discovery, stereochemistry or atom-name lookup is part of this operation.
For real x, $asinh(x)=\log(x+\sqrt{1+x^2})$ . A stable value formula uses $sign(x)\log(|x|+\sqrt{1+x^2})$ ; handle zero so its derivative remains 1.

Why this matters

A field’s vocabulary defines its meaning: an unknown observed value needs an explicit category, while a missing field contributes no observation.

Starter code

You write the body of this function in the BioTorch workspace. Your code runs in an isolated PyTorch sandbox; Run checks the public examples and Submit grades the full test suite, including gradients.

import torch


def categorical_feature_encoding(
    categories: torch.Tensor,  # (..., F)
    choices: torch.Tensor,  # (F, K)
    continuous: torch.Tensor,  # (..., D)
    continuous_transform: torch.Tensor,  # (D,)
    continuous_scale: torch.Tensor,  # (D,)
    parameters: dict,
    masks: dict,
    unknown_policy: str = 'bucket',
) -> dict[str, torch.Tensor]:
    """Return dict[str, torch.Tensor]:
    
    one_hot: (..., F, K+1).
    features: (..., F* (K+1)+D).
    embedding: (..., O)."""
    # Implement the operation. Inputs and test setup are supplied by BioTorch.
    raise NotImplementedError("Implement categorical_feature_encoding")

Ready to implement it?

3 progressive hints, a solution walkthrough, and grading that checks values, shapes and gradients are waiting in the workspace.

Write and run your solution Sign in

Free account with email, Google or GitHub. Your code runs in an isolated PyTorch sandbox.

Sources and further reading

A primitive exercise, not an end-to-end reproduction of the named models. No affiliation or endorsement is implied.