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.
    """
    if self._mode == 'train' or self._multiply_by_controller_output:
      quant_mask, mask, x = x
    else:
      quant_mask, x = x
    original_quant_mask = quant_mask

    w1, w2, b2 = self.weights

    if self._mode == 'predict':
      w1 = jnp.transpose(w1, (1, 2, 0))  # dm, d1, d2 -> d1, d2, dm
      w2 = jnp.transpose(w2, (1, 0, 2))  # d2, d1, dm -> d1, d2, dm
    x_shape = x.shape
    x = jnp.reshape(x, [-1, x_shape[-1]])  # Easier to operate on flattened x.

    if self._mode == 'train':
      # Tricks from Section 2.1 in https://arxiv.org/abs/1801.09797
      quant_mask = tl.one_hot(quant_mask, self._n_elements_in_block)
      quant_mask = fastmath.stop_gradient(quant_mask)
      quant_mask += mask - fastmath.stop_gradient(mask)  # straight-through
      # We will sometimes (quant_prob of the batches) use the soft-mask instead
      # of the quantized mask to improve training stability (see paper above).
      select = fastmath.random.uniform(self.rng, (), jnp.float32, 0.0, 1.0)
      quant_mask = jnp.where(select < self._quant_prob, quant_mask, mask)

      # In training, run full matmul to get benefits from the above tricks.
      mid = jnp.einsum('bd,dxy->bxy', x, w1) * quant_mask
      relu = jnp.where(mid <= 0, jnp.zeros_like(mid), mid)
      if self._multiply_by_controller_output:
        # We multiply only for quantized decisions, since for non-quantized
        # decisions we've already multiplied the output.
        mask_mult = jnp.where(select < self._quant_prob,
                              mask, jnp.ones_like(mask))
        # Stop-gradient is here, because we already have a pass-through gradient
        # (for quantized decisions).
        mask_mult = fastmath.stop_gradient(mask_mult)
        relu = relu * mask_mult
      res = jnp.einsum('bxy,yxd->bd', relu, w2) + b2
    elif self._mode == 'predict':
      # This implementation mimicks inference. It's not efficient for large
      # size of joint_batch, but at inference that will be 1 most of the time.
      # Shapes:
      # quant_mask is [joint_batch, self._d1]
      # w1 is [d_model, self._d1, self._d2]
      # we'll index w1 with advanced numpy indexing, first range over
      # self._d1 times the batch size, second range being quant_mask
      batch_size = quant_mask.shape[0]
      idx1 = jnp.array([jnp.arange(self._d1)] * batch_size)
      # flatten indices and select from w1
      idx1 = jnp.reshape(idx1, [-1])
      idx2 = jnp.reshape(quant_mask, [-1])
      w = w1[idx1, idx2, :]  # now we have per-element weights with batch dim
      w = jnp.reshape(w, [batch_size, self._d1, -1])
      mid = jnp.einsum('ai,aji->aj', x, w)
      relu = jnp.where(mid <= 0, jnp.zeros_like(mid), mid)
      if self._multiply_by_controller_output:
        mask_mult = jnp.take_along_axis(mask, quant_mask[..., None], -1)[..., 0]
        relu = relu * mask_mult
      # w2 is [self._d1, self._d2, d_model]
      v = w2[idx1, idx2, :]
      v = jnp.reshape(v, [batch_size, self._d1, -1])
      res = jnp.einsum('ai,aij->aj', relu, v) + b2
    else:
      quant_mask = tl.one_hot(quant_mask, self._n_elements_in_block)
      mid = jnp.einsum('bd,dxy->bxy', x, w1) * quant_mask
      relu = jnp.where(mid <= 0, jnp.zeros_like(mid), mid)
      if self._multiply_by_controller_output:
        relu = relu * mask
      res = jnp.einsum('bxy,yxd->bd', relu, w2) + b2

    return original_quant_mask, jnp.reshape(res, x_shape)
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