replay-rec 0.20.3__py3-none-any.whl → 0.21.0__py3-none-any.whl

replay/models/nn/sequential/sasrec/model.py

@@ -3,10 +3,15 @@ import contextlib
 from typing import Any, Optional, Union, cast
 
 import torch
+from typing_extensions import deprecated
 
 from replay.data.nn import TensorMap, TensorSchema
 
 
+@deprecated(
+    "`SasRecModel` class is deprecated. Use `replay.nn.sequential.SasRec` instead.",
+    stacklevel=2,
+)
 class SasRecModel(torch.nn.Module):
     """
     SasRec model
@@ -110,7 +115,7 @@ class SasRecModel(torch.nn.Module):
     ) -> torch.Tensor:
         """
         :param feature_tensor: Batch of features.
-        :param padding_mask: Padding mask where 0 - <PAD>, 1 otherwise.
+        :param padding_mask: Padding mask where 0 - ``<PAD>``, 1 - otherwise.
 
         :returns: Calculated scores.
         """
@@ -127,11 +132,11 @@ class SasRecModel(torch.nn.Module):
     ) -> torch.Tensor:
         """
         :param feature_tensor: Batch of features.
-        :param padding_mask: Padding mask where 0 - <PAD>, 1 otherwise.
-        :param candidates_to_score: Item ids to calculate scores.
-            if `None` predicts for all items
+        :param padding_mask: Padding mask where 0 - ``<PAD>``, 1 - otherwise.
+        :param candidates_to_score: Item ids to calculate scores.\n
+            If ``None`` then predicts for all items. Default: ``None``.
 
-        :returns: Prediction among canditates_to_score items.
+        :returns: Prediction among ``canditates_to_score`` items.
         """
         # final_emb: [B x E]
         final_emb = self.get_query_embeddings(feature_tensor, padding_mask)
@@ -145,7 +150,7 @@ class SasRecModel(torch.nn.Module):
     ):
         """
         :param feature_tensor: Batch of features.
-        :param padding_mask: Padding mask where 0 - <PAD>, 1 otherwise.
+        :param padding_mask: Padding mask where 0 - ``<PAD>``, 1 - otherwise.
 
         :returns: Query embeddings.
         """
@@ -158,7 +163,7 @@ class SasRecModel(torch.nn.Module):
     ) -> torch.Tensor:
         """
         :param feature_tensor: Batch of features.
-        :param padding_mask: Padding mask where 0 - <PAD>, 1 otherwise.
+        :param padding_mask: Padding mask where 0 - ``<PAD>``, 1 - otherwise.
 
         :returns: Output embeddings.
         """
@@ -176,9 +181,9 @@ class SasRecModel(torch.nn.Module):
 
     def get_logits(self, out_embeddings: torch.Tensor, item_ids: Optional[torch.LongTensor] = None) -> torch.Tensor:
         """
-        Apply head to output embeddings of `forward_step`.
+        Apply head to output embeddings of ``forward_step``.
 
-        :param out_embeddings: Embeddings after `forward step`.
+        :param out_embeddings: Embeddings after ``forward step``.
         :param item_ids: Item ids to calculate scores.
             Default: ``None``.
 

replay/nn/__init__.py

@@ -0,0 +1,8 @@
+from replay.utils import TORCH_AVAILABLE
+
+if not TORCH_AVAILABLE:
+    msg = (
+        "The replay.nn module is unavailable. "
+        "To use the functionality from this module, please install ``torch`` and ``lightning``."
+    )
+    raise ImportError(msg)

replay/nn/agg.py

@@ -0,0 +1,109 @@
+import contextlib
+from typing import Protocol
+
+import torch
+
+from replay.data.nn.schema import TensorMap
+
+
+class AggregatorProto(Protocol):
+    """Class-protocol for working with embedding aggregation functions"""
+
+    def forward(
+        self,
+        feature_tensors: TensorMap,
+    ) -> torch.Tensor: ...
+
+    @property
+    def embedding_dim(self) -> int: ...
+
+    def reset_parameters(self) -> None: ...
+
+
+class SumAggregator(torch.nn.Module):
+    """
+    The class summarizes the incoming embeddings.
+    Note that for successful aggregation, the dimensions of all embeddings must match.
+    """
+
+    def __init__(self, embedding_dim: int) -> None:
+        """
+        :param embedding_dim: The last dimension of incoming and outcoming embeddings.
+        """
+        super().__init__()
+        self._embedding_dim = embedding_dim
+
+    @property
+    def embedding_dim(self) -> int:
+        """The dimension of the output embedding"""
+        return self._embedding_dim
+
+    def reset_parameters(self) -> None:
+        pass
+
+    def forward(self, feature_tensors: TensorMap) -> torch.Tensor:
+        """
+        :param feature_tensors: a dictionary of tensors to sum up.
+            The dimensions of all tensors in the dictionary must match.
+
+        :returns: torch.Tensor. The last dimension of the tensor is ``embedding_dim``.
+        """
+        out = sum(feature_tensors.values())
+        assert out.size(-1) == self.embedding_dim
+        return out
+
+
+class ConcatAggregator(torch.nn.Module):
+    """
+    The class concatenates incoming embeddings by the last dimension.
+
+    If you need to concatenate several embeddings,
+    then a linear layer will be applied to get the last dimension equal to ``embedding_dim``.\n
+    If only one embedding comes to the input, then its last dimension is expected to be equal to ``embedding_dim``.
+    """
+
+    def __init__(
+        self,
+        input_embedding_dims: list[int],
+        output_embedding_dim: int,
+    ) -> None:
+        """
+        :param input_embedding_dims: Dimensions of incoming embeddings.
+        :param output_embedding_dim: The dimension of the output embedding after concatenation.
+        """
+        super().__init__()
+        self._embedding_dim = output_embedding_dim
+        embedding_concat_size = sum(input_embedding_dims)
+        self.feat_projection = None
+        if len(input_embedding_dims) > 1:
+            self.feat_projection = torch.nn.Linear(embedding_concat_size, self.embedding_dim)
+        elif embedding_concat_size != self.embedding_dim:
+            msg = f"Input embedding dim is not equal to embedding_dim ({embedding_concat_size} != {self.embedding_dim})"
+            raise ValueError(msg)
+
+    @property
+    def embedding_dim(self) -> int:
+        """The dimension of the output embedding"""
+        return self._embedding_dim
+
+    def reset_parameters(self) -> None:
+        for _, param in self.named_parameters():
+            with contextlib.suppress(ValueError):
+                torch.nn.init.xavier_normal_(param.data)
+
+    def forward(self, feature_tensors: TensorMap) -> torch.Tensor:
+        """
+        To ensure the deterministic nature of the result,
+        the embeddings are concatenated in the ascending order of the keys in the dictionary.
+
+        :param feature_tensors: a dictionary of tensors to concatenate.
+
+        :returns: The last dimension of the tensor is ``embedding_dim``.
+        """
+        # To maintain determinism, we concatenate the tensors in sorted order by names.
+        sorted_names = sorted(feature_tensors.keys())
+        out = torch.cat([feature_tensors[name] for name in sorted_names], dim=-1)
+        if self.feat_projection is not None:
+            out = self.feat_projection(out)
+        assert out.size(-1) == self.embedding_dim
+        return out

replay/nn/attention.py

@@ -0,0 +1,158 @@
+import contextlib
+import math
+from typing import Optional
+
+import torch
+
+
+class MultiHeadDifferentialAttention(torch.nn.Module):
+    """
+    Multi-Head Differential Attention Mechanism.
+    Replaces the conventional softmax attention with a differential attention.
+    Incorporattes a causal mask (if other not specified) to ensure autoregressive behavior.
+
+    Source paper: https://arxiv.org/pdf/2410.05258
+    """
+
+    def __init__(
+        self,
+        embedding_dim: int,
+        num_heads: int,
+        lambda_init: float,
+        bias: bool = False,
+        kdim: Optional[int] = None,
+        vdim: Optional[int] = None,
+    ):
+        """
+        :param embedding_dim: Total dimension of the model. Must be divisible by ``num_heads``.
+        :param num_heads: Number of parallel attention heads.
+        :param lambda_init: Initial value for lambda.
+        :param bias: If specified, adds bias to input / output projection layers. Default: ``False``.
+        :param kdim: Total number of features for keys. Default: ``None`` (uses ``kdim=embedding_dim``).
+        :param vdim: Total number of features for values. Default: ``None`` (uses ``vdim=embedding_dim``).
+        """
+        super().__init__()
+        kdim = kdim or embedding_dim
+        vdim = vdim or embedding_dim
+        assert kdim % num_heads == 0, "Query/Key embedding dim is not divisible by num_heads"
+        assert vdim % num_heads == 0, "Value embedding dim is not divisible by num_heads"
+        self.qk_e_head = kdim // num_heads
+        self.v_e_head = vdim // num_heads
+        self.num_heads = num_heads
+
+        # Linear projections for queries, keys, and values
+        # Project to 2 * d_head per head for differential attention
+        self.W_q = torch.nn.Linear(embedding_dim, 2 * self.qk_e_head * num_heads, bias=bias)
+        self.W_k = torch.nn.Linear(embedding_dim, 2 * self.qk_e_head * num_heads, bias=bias)
+        self.W_v = torch.nn.Linear(embedding_dim, self.v_e_head * num_heads, bias=bias)
+        self.W_o = torch.nn.Linear(self.v_e_head * num_heads, embedding_dim, bias=bias)
+
+        # Learnable parameters for lambda reparameterization
+        self.lambda_q1 = torch.nn.Parameter(torch.randn(num_heads, self.qk_e_head))
+        self.lambda_k1 = torch.nn.Parameter(torch.randn(num_heads, self.qk_e_head))
+        self.lambda_q2 = torch.nn.Parameter(torch.randn(num_heads, self.qk_e_head))
+        self.lambda_k2 = torch.nn.Parameter(torch.randn(num_heads, self.qk_e_head))
+        self.register_buffer("scaling", torch.asarray(1 / math.sqrt(self.qk_e_head), dtype=torch.float32))
+
+        self.lambda_init = lambda_init
+
+        # Scale parameter for RMSNorm
+        self.rms_scale = torch.nn.Parameter(torch.ones(self.v_e_head))
+        self.eps = 1e-5  # Epsilon for numerical stability
+
+    def reset_parameters(self) -> None:
+        for _, param in self.named_parameters():
+            with contextlib.suppress(ValueError):
+                torch.nn.init.xavier_normal_(param.data)
+
+    def forward(
+        self,
+        query: torch.Tensor,
+        key: torch.Tensor,
+        value: torch.Tensor,
+        attn_mask: torch.FloatTensor,
+    ) -> torch.Tensor:
+        """
+        Forward pass for Multi-Head Differential Attention.
+
+        :param query: Query sequence of shape ``(batch_size, sequence_length, embedding_dim)``.
+        :param key: Key sequence of shape ``(batch_size, sequence_length, embedding_dim)``.
+        :param value: Value sequence of shape ``(batch_size, sequence_length, embedding_dim)``.
+        :param attn_mask: attention mask, where ``-inf`` for ``PAD``, ``0`` - otherwise.\n
+            Possible shapes:\n
+            1. ``(batch_size * num_heads, sequence_length, sequence_length)``
+            2. ``(batch_size, num_heads, sequence_length, sequence_length)``
+        :returns: torch.Tensor: Output tensor after applying differential attention.
+        """
+        batch_size, seq_len, _ = value.shape
+
+        # Project inputs to queries, keys, and values
+        query = self.W_q(query)  # Shape: (batch_size, seq_len, 2 * num_heads * qk_e_head)
+        key = self.W_k(key)  # Shape: (batch_size, seq_len, 2 * num_heads * qk_e_head)
+        value = self.W_v(value)  # Shape: (batch_size, seq_len, num_heads * v_e_head)
+
+        # Reshape and permute for multi-head attention
+        # New shape: (batch_size, num_heads, sequence_length, 2 * qk_e_head or v_e_head)
+        query = query.view(batch_size, seq_len, self.num_heads, 2 * self.qk_e_head).transpose(1, 2)
+        key = key.view(batch_size, seq_len, self.num_heads, 2 * self.qk_e_head).transpose(1, 2)
+        value = value.view(batch_size, seq_len, self.num_heads, self.v_e_head).transpose(1, 2)
+
+        # Split query and key into query1, query2 and key1, key2
+        query1, query2 = query.chunk(2, dim=-1)  # Each of shape: (batch_size, num_heads, seq_len, d_head)
+        key1, key2 = key.chunk(2, dim=-1)  # Each of shape: (batch_size, num_heads, seq_len, d_head)
+
+        # Compute lambda using reparameterization
+        # Compute dot products for each head
+        # Shape of lambda_val: (num_heads,)
+        lambda_q1_dot_k1 = torch.sum(self.lambda_q1 * self.lambda_k1, dim=-1).float()  # (num_heads,)
+        lambda_q2_dot_k2 = torch.sum(self.lambda_q2 * self.lambda_k2, dim=-1).float()  # (num_heads,)
+        lambda_val = torch.exp(lambda_q1_dot_k1) - torch.exp(lambda_q2_dot_k2) + self.lambda_init  # (num_heads,)
+
+        # Expand lambda_val to match attention dimensions (batch_size, num_heads, 1, 1)
+        lambda_val = lambda_val.unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
+
+        # Reshape attn_mask from 3D  to 4D
+        if len(attn_mask.shape) == 3:
+            attn_mask = attn_mask.reshape(attn_mask.shape[0] // self.num_heads, self.num_heads, *attn_mask.shape[1:])
+
+        # check shapes
+        assert attn_mask.dim() == 4
+        assert attn_mask.size() == (batch_size, self.num_heads, seq_len, seq_len)
+
+        # Compute attention scores
+        attention_scores1 = torch.matmul(query1, key1.transpose(-2, -1)) * self.get_buffer(
+            "scaling"
+        )  # (batch_size, num_heads, seq_len, seq_len)
+        attention_scores2 = torch.matmul(query2, key2.transpose(-2, -1)) * self.get_buffer(
+            "scaling"
+        )  # (batch_size, num_heads, seq_len, seq_len)
+
+        # Apply the causal mask
+        attention_scores1 = attention_scores1 + attn_mask  # Mask out future positions
+        attention_scores2 = attention_scores2 + attn_mask  # Mask out future positions
+
+        # Apply softmax to get attention weights
+        attention1 = torch.nn.functional.softmax(attention_scores1, dim=-1)  # (batch_size, num_heads, seq_len, seq_len)
+        attention2 = torch.nn.functional.softmax(attention_scores2, dim=-1)
+        attention = attention1 - lambda_val * attention2
+
+        # Apply attention weights to values
+        output = torch.matmul(attention, value)  # (batch_size, num_heads, seq_len, v_e_head)
+
+        # Normalize each head independently using RMSNorm
+        # Compute RMSNorm
+        rms_norm = torch.sqrt(output.pow(2).mean(dim=-1, keepdim=True) + self.eps)  # (batch_size*num_heads, seq_len, 1)
+        output_normalized = (output / rms_norm) * self.rms_scale  # (batch*num_heads, seq_len, v_e_head)
+
+        # Scale the normalized output
+        output_normalized = output_normalized * (1 - self.lambda_init)  # Scalar scaling
+
+        # Concatenate all heads
+        # New shape: (batch_size, seq_len, num_heads * v_e_head)
+        output_concat = (
+            output_normalized.transpose(1, 2).contiguous().view(batch_size, seq_len, self.num_heads * self.v_e_head)
+        )
+
+        # Final linear projection
+        output_projection = self.W_o(output_concat)  # (batch_size, seq_len, embedding_dim)
+        return output_projection

replay/nn/embedding.py

@@ -0,0 +1,283 @@
+import contextlib
+import warnings
+from collections.abc import Sequence
+from typing import Literal, Optional, Union
+
+import torch
+
+from replay.data.nn.schema import TensorFeatureInfo, TensorMap, TensorSchema
+
+
+class SequenceEmbedding(torch.nn.Module):
+    """
+    The embedding generation class for all types of features given into the sequential models.
+
+    The embedding size for each feature will be taken from ``TensorSchema`` (from field named ``embedding_dim``).
+    For numerical features, it is expected that the last dimension of the tensor will be equal
+    to ``tensor_dim`` field in ``TensorSchema``.
+
+    Keep in mind that the first dimension of the every categorical embedding (the size of embedding table)
+    will equal to the ``cardinality`` + 1. This is necessary to take into account the padding value.
+    """
+
+    def __init__(
+        self,
+        schema: TensorSchema,
+        excluded_features: Optional[list[str]] = None,
+        categorical_list_feature_aggregation_method: Literal["sum", "mean", "max"] = "sum",
+    ):
+        """
+        :param schema: TensorSchema containing meta information about all the features
+            for which you need to generate an embedding.
+        :param excluded_features: A list containing the names of features
+            for which you do not need to generate an embedding.
+            Fragments from this list are expected to be contained in ``schema``.
+            Default: ``None``.
+        :param categorical_list_feature_aggregation_method: Mode to aggregate tokens
+            in token item representation (categorical list only).
+            Default: ``"sum"``.
+        """
+        super().__init__()
+        self.excluded_features = excluded_features or []
+        feature_embedders = {}
+
+        for feature_name, tensor_info in schema.items():
+            if feature_name in self.excluded_features:
+                continue
+            if not tensor_info.is_seq:
+                msg = f"Non-sequential features is not yet supported. Got {feature_name}"
+                raise NotImplementedError(msg)
+            if tensor_info.is_cat:
+                feature_embedders[feature_name] = CategoricalEmbedding(
+                    tensor_info,
+                    categorical_list_feature_aggregation_method,
+                )
+            else:
+                feature_embedders[feature_name] = NumericalEmbedding(tensor_info)
+
+        self.feature_names = list(feature_embedders.keys())
+        if not feature_embedders:
+            msg = "Expected to have at least one feature name to generate embedding."
+            raise ValueError(msg)
+        self.feature_embedders: dict[str, Union[CategoricalEmbedding, NumericalEmbedding]] = torch.nn.ModuleDict(
+            feature_embedders
+        )
+        self._item_feature_name = schema.item_id_feature_name
+
+    def reset_parameters(self) -> None:
+        for _, param in self.named_parameters():
+            with contextlib.suppress(ValueError):
+                torch.nn.init.xavier_normal_(param.data)
+
+    def forward(self, feature_tensor: TensorMap, feature_names: Optional[Sequence[str]] = None) -> TensorMap:
+        """
+        :param feature_tensor: a dictionary of tensors to generate embedding.
+            It is expected that the keys from this dictionary match the names of the features in the given ``schema``.
+        :param feature_names: A custom list of features for which embeddings need to be generated.
+            It is expected that the values from this list match the names of the features in the given ``schema``.\n
+            Default: ``None``. This means that the names of the features from the ``schema`` will be used.
+
+        :returns: a dictionary with tensors that contains embeddings.
+        """
+        return {
+            feature_name: self.feature_embedders[feature_name](feature_tensor[feature_name])
+            for feature_name in (feature_names or self.feature_names)
+        }
+
+    @property
+    def embeddings_dim(self) -> dict[str, int]:
+        """
+        Returns the embedding dimensions for each of the features in the `schema`.
+        """
+        return {name: emb.embedding_dim for name, emb in self.feature_embedders.items()}
+
+    def get_item_weights(self, indices: Optional[torch.LongTensor] = None) -> torch.Tensor:
+        """
+        Getting the embedding weights for a feature that matches the item id feature
+        with the name specified in the ``schema``.
+        It is expected that embeddings for this feature will definitely exist.
+        **Note**: the row corresponding to the padding will be excluded from the returned weights.
+        This logic will work if given ``indices`` is ``None``.
+
+        :param indices: Items indices.
+        :returns: Embeddings for specific items.
+        """
+        if indices is None:
+            return self.feature_embedders[self._item_feature_name].weight
+        return self.feature_embedders[self._item_feature_name](indices)
+
+
+class CategoricalEmbedding(torch.nn.Module):
+    """
+    The embedding generation class for categorical features.
+    It supports working with single features for each event in sequence, as well as several (categorical list).
+
+    When using this class, keep in mind that
+    the first dimension of the embedding (the size of embedding table) will equal to the ``cardinality`` + 1.
+    This is necessary to take into account the padding value.
+    """
+
+    def __init__(
+        self,
+        feature_info: TensorFeatureInfo,
+        categorical_list_feature_aggregation_method: Literal["sum", "mean", "max"] = "sum",
+    ) -> None:
+        """
+        :param feature_info: Meta information about the feature.
+        :param categorical_list_feature_aggregation_method: Mode to aggregate tokens
+            in token item representation (categorical list only). One of {`sum`, `mean`, `max`}
+            Default: ``"sum"``.
+        """
+        super().__init__()
+        assert feature_info.cardinality
+        assert feature_info.embedding_dim
+
+        self._expect_padding_value_setted = True
+        if feature_info.cardinality != feature_info.padding_value:
+            self._expect_padding_value_setted = False
+            msg = (
+                f"The padding value={feature_info.padding_value} is set for the feature={feature_info.name}. "
+                f"The expected padding value for this feature should be {feature_info.cardinality}. "
+                "Keep this in mind when getting the weights via the `weight` property, "
+                "because the weights are returned there without padding row. "
+                "Therefore, during the IDs scores generating, "
+                "all the IDs that greater than the padding value should be increased by 1."
+            )
+            warnings.warn(msg, stacklevel=2)
+
+        if feature_info.is_list:
+            self.emb = torch.nn.EmbeddingBag(
+                feature_info.cardinality + 1,
+                feature_info.embedding_dim,
+                padding_idx=feature_info.padding_value,
+                mode=categorical_list_feature_aggregation_method,
+            )
+            self._get_embeddings = self._get_cat_list_embeddings
+        else:
+            self.emb = torch.nn.Embedding(
+                feature_info.cardinality + 1,
+                feature_info.embedding_dim,
+                padding_idx=feature_info.padding_value,
+            )
+            self._get_embeddings = self._get_cat_embeddings
+
+    @property
+    def weight(self) -> torch.Tensor:
+        """
+        Returns the weights of the embedding layer,
+        excluding the row that corresponds to the padding.
+        """
+        if not self._expect_padding_value_setted:
+            msg = (
+                "The weights are returned there do not contain padding row. "
+                "Therefore, during the IDs scores generating, "
+                "all the IDs that greater than the padding value should be increased by 1."
+            )
+            warnings.warn(msg, stacklevel=2)
+
+        mask_without_padding = torch.ones(
+            size=(self.emb.weight.size(0),),
+            dtype=torch.bool,
+            device=self.emb.weight.device,
+        )
+        mask_without_padding[self.emb.padding_idx].zero_()
+        return self.emb.weight[mask_without_padding]
+
+    def forward(self, indices: torch.LongTensor) -> torch.Tensor:
+        """
+        :param indices: Items indices.
+
+        :returns: Embeddings for specific items.
+        """
+        return self._get_embeddings(indices)
+
+    @property
+    def embedding_dim(self) -> int:
+        """Embedding dimension after applying the layer"""
+        return self.emb.embedding_dim
+
+    def _get_cat_embeddings(self, indices: torch.LongTensor) -> torch.Tensor:
+        """
+        :param indices: Items indices.
+
+        :returns: Embeddings for specific items.
+        """
+        return self.emb(indices)
+
+    def _get_cat_list_embeddings(self, indices: torch.LongTensor) -> torch.Tensor:
+        """
+        :param indices: Items indices.
+
+        :returns: Embeddings for specific items.
+        """
+        assert indices.dim() >= 2
+        if indices.dim() == 2:
+            embeddings: torch.Tensor = self.emb(indices)
+        else:
+            source_size = indices.size()
+            indices = indices.view(-1, source_size[-1])
+            embeddings = self.emb(indices)
+            embeddings = embeddings.view(*source_size[:-1], -1)
+        return embeddings
+
+
+class NumericalEmbedding(torch.nn.Module):
+    """
+    The embedding generation class for numerical features.
+    It supports working with single features for each event in sequence, as well as several (numerical list).
+
+    **Note**: if the ``embedding_dim`` field in ``TensorFeatureInfo`` for an incoming feature matches its last dimension
+    (``tensor_dim`` field in ``TensorFeatureInfo``), then transformation will not be applied.
+    """
+
+    def __init__(self, feature_info: TensorFeatureInfo) -> None:
+        """
+        :param feature_info: Meta information about the feature.
+        """
+        super().__init__()
+        assert feature_info.tensor_dim
+        assert feature_info.embedding_dim
+        self._tensor_dim = feature_info.tensor_dim
+        self._embedding_dim = feature_info.embedding_dim
+        self.linear = torch.nn.Linear(feature_info.tensor_dim, self.embedding_dim)
+
+        if feature_info.is_list:
+            if self.embedding_dim == feature_info.tensor_dim:
+                torch.nn.init.eye_(self.linear.weight.data)
+                torch.nn.init.zeros_(self.linear.bias.data)
+
+                self.linear.weight.requires_grad = False
+                self.linear.bias.requires_grad = False
+        else:
+            assert feature_info.tensor_dim == 1
+            self.linear = torch.nn.Linear(feature_info.tensor_dim, self.embedding_dim)
+
+    @property
+    def weight(self) -> torch.Tensor:
+        """
+        Returns the weight of the applied layer.
+        If ``embedding_dim`` matches ``tensor_dim``, then the identity matrix will be returned.
+        """
+        return self.linear.weight
+
+    def forward(self, values: torch.FloatTensor) -> torch.Tensor:
+        """
+        Numerical embedding forward pass.\n
+        **Note**: if the ``embedding_dim`` for an incoming feature matches its last dimension (``tensor_dim``),
+        then transformation will not be applied.
+
+        :param values: feature values.
+        :returns: Embeddings for specific items.
+        """
+        if values.dim() <= 2 and self._tensor_dim == 1:
+            values = values.unsqueeze(-1).contiguous()
+
+        assert values.size(-1) == self._tensor_dim
+        if self._tensor_dim != self.embedding_dim:
+            return self.linear(values)
+        return values
+
+    @property
+    def embedding_dim(self) -> int:
+        """Embedding dimension after applying the layer"""
+        return self._embedding_dim