"""Apply a linear layer to sparse feature rows."""

import torch.nn as nn
from . import SparseTensor


class SparseLinear(nn.Linear):
    def forward(self, value: SparseTensor) -> SparseTensor:
        features = value.feats.to(self.weight.dtype)
        return value.replace(super().forward(features))