def forward(self, x):
    """Executes this layer as part of a forward pass through the model.

    Args:
      x: Tensor of same shape and dtype as the input signature used to
          initialize this layer.

    Returns:
      Tensor of same shape and dtype as the input.
    """
    x, recalled_quant_mask = x
    m1, m2, mb = self.weights

    x_shape = x.shape
    x = jnp.reshape(x, [-1, x_shape[-1]])  # Easier to operate on flattened x.

    # Q: should we add bias and/or put relu after the low-rank m1 dot?
    # Replacing multiplication and reshape by this einsum brings training speed
    # improvement (see also reshape in initialization).
    mask_logits = jnp.einsum('bd,dl,lxy->bxy', x, m1, m2) + mb

    if self._also_return_nondiscrete_output:
      # Softmax.
      mask_logsumexp = fastmath.logsumexp(mask_logits, axis=-1, keepdims=True)
      log_mask = mask_logits - mask_logsumexp
      mask = jnp.exp(log_mask)
      # Gumbel-softmax with straight-through discretization.
      if self._temperature == 0.0:
        quant_mask = jnp.argmax(log_mask, axis=-1)
      else:
        u = fastmath.random.uniform(self.rng, mask.shape, jnp.float32, 1e-6,
                                    1.0 - 1e-6)
        g = -jnp.log(-jnp.log(u))
        quant_mask = jnp.argmax(log_mask + g * self._temperature, axis=-1)
    else:
      quant_mask = jnp.argmax(mask_logits, axis=-1)

    if self._mode == 'train':
      # We use recalled_quant_mask if it's different than -1; otherwise
      # we use a quant_mask which we have just computed.
      quant_mask = jnp.where(recalled_quant_mask == -1,
                             quant_mask, recalled_quant_mask)

    if self._also_return_nondiscrete_output:
      return quant_mask, mask
    else:
      return quant_mask
Executes this layer as part of a forward pass through the model.

    Args:
      x: Tensor of same shape and dtype as the input signature used to
          initialize this layer.

    Returns:
      Tensor of same shape and dtype as the input.
    
trax/layers/research/sparsity.py