encoder.py

from modules.rqvae.normalize import L2NormalizationLayer, RMSNorm
from typing import List
from torch import nn
from torch import Tensor


class MLP(nn.Module):
    def __init__(
        self,
        input_dim: int,
        hidden_dims: List[int],
        out_dim: int,
        dropout: float = 0.0,
        normalize: bool = False,
        normalize_type: str = 'rms'
    ) -> None:
        super().__init__()

        self.input_dim = input_dim
        self.hidden_dims = hidden_dims
        self.out_dim = out_dim
        self.dropout = dropout

        dims = [self.input_dim] + self.hidden_dims + [self.out_dim]
        
        self.mlp = nn.Sequential()
        for i, (in_d, out_d) in enumerate(zip(dims[:-1], dims[1:])):
            self.mlp.append(nn.Linear(in_d, out_d, bias=False))
            if i != len(dims)-2:
                self.mlp.append(nn.SiLU())
                if dropout != 0:
                    self.mlp.append(nn.Dropout(dropout))
        if normalize:
            if normalize_type == 'l2':
                self.mlp.append(L2NormalizationLayer())
            elif normalize_type == 'rms':
                self.mlp.append(RMSNorm())
            else:
                raise NotImplementedError()
        else:
            self.mlp.append(nn.Identity())

    def forward(self, x: Tensor) -> Tensor:
        assert x.shape[-1] == self.input_dim, f"Invalid input dim: Expected {self.input_dim}, found {x.shape[-1]}"
        return self.mlp(x)