| from flax import linen as nn |
| import jax |
|
|
|
|
| class SequenceLayer(nn.Module): |
| """ Defines a single S5 layer, with S5 SSM, nonlinearity, |
| dropout, batch/layer norm, etc. |
| Args: |
| ssm (nn.Module): the SSM to be used (i.e. S5 ssm) |
| dropout (float32): dropout rate |
| d_model (int32): this is the feature size of the layer inputs and outputs |
| we usually refer to this size as H |
| activation (string): Type of activation function to use |
| training (bool): whether in training mode or not |
| prenorm (bool): apply prenorm if true or postnorm if false |
| batchnorm (bool): apply batchnorm if true or layernorm if false |
| bn_momentum (float32): the batchnorm momentum if batchnorm is used |
| step_rescale (float32): allows for uniformly changing the timescale parameter, |
| e.g. after training on a different resolution for |
| the speech commands benchmark |
| """ |
| ssm: nn.Module |
| dropout: float |
| d_model: int |
| activation: str = "gelu" |
| training: bool = True |
| prenorm: bool = False |
| batchnorm: bool = False |
| bn_momentum: float = 0.90 |
| step_rescale: float = 1.0 |
|
|
| def setup(self): |
| """Initializes the ssm, batch/layer norm and dropout |
| """ |
| self.seq = self.ssm(step_rescale=self.step_rescale) |
|
|
| if self.activation in ["full_glu"]: |
| self.out1 = nn.Dense(self.d_model) |
| self.out2 = nn.Dense(self.d_model) |
| elif self.activation in ["half_glu1", "half_glu2"]: |
| self.out2 = nn.Dense(self.d_model) |
|
|
| if self.batchnorm: |
| self.norm = nn.BatchNorm(use_running_average=not self.training, |
| momentum=self.bn_momentum, axis_name='batch') |
| else: |
| self.norm = nn.LayerNorm() |
|
|
| self.drop = nn.Dropout( |
| self.dropout, |
| broadcast_dims=[0], |
| deterministic=not self.training, |
| ) |
|
|
| def __call__(self, x): |
| """ |
| Compute the LxH output of S5 layer given an LxH input. |
| Args: |
| x (float32): input sequence (L, d_model) |
| Returns: |
| output sequence (float32): (L, d_model) |
| """ |
| skip = x |
| if self.prenorm: |
| x = self.norm(x) |
| x = self.seq(x) |
|
|
| if self.activation in ["full_glu"]: |
| x = self.drop(nn.gelu(x)) |
| x = self.out1(x) * jax.nn.sigmoid(self.out2(x)) |
| x = self.drop(x) |
| elif self.activation in ["half_glu1"]: |
| x = self.drop(nn.gelu(x)) |
| x = x * jax.nn.sigmoid(self.out2(x)) |
| x = self.drop(x) |
| elif self.activation in ["half_glu2"]: |
| |
| x1 = self.drop(nn.gelu(x)) |
| x = x * jax.nn.sigmoid(self.out2(x1)) |
| x = self.drop(x) |
| elif self.activation in ["gelu"]: |
| x = self.drop(nn.gelu(x)) |
| else: |
| raise NotImplementedError( |
| "Activation: {} not implemented".format(self.activation)) |
|
|
| x = skip + x |
| if not self.prenorm: |
| x = self.norm(x) |
| return x |
|
|
