torch_frame.nn.encoder.MultiCategoricalEmbeddingEncoder

class MultiCategoricalEmbeddingEncoder(out_channels: Optional[int] = None, stats_list: Optional[list[dict[torch_frame.data.stats.StatType, Any]]] = None, stype: Optional[stype] = None, post_module: Optional[Module] = None, na_strategy: Optional[NAStrategy] = None, mode: str = 'mean')

Bases: StypeEncoder

An embedding look-up based encoder for multi_categorical features. It applies torch.nn.EmbeddingBag for each categorical feature and concatenates the output embeddings.

Parameters:

mode (str) – “sum”, “mean” or “max”. Specifies the way to reduce the bag. (default: mean)

reset_parameters() → None

Initialize the parameters of post_module.

encode_forward(feat: MultiNestedTensor, col_names: Optional[list[str]] = None) → Tensor

The main forward function. Maps input feat from TensorFrame (shape [batch_size, num_cols]) into output x of shape [batch_size, num_cols, out_channels].