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

    Args:
      x: Tensor of activations.

    Returns:
      Tensor of same shape and dtype as the input.
    """
    if self._mode != 'train':
      return x
    state, rng = self.state, self.rng
    rate = self._initial_rate
    if isinstance(state, dict) and self._name in state:
      rate = state[self._name]
    if rate == 0.0:
      return x
    mask_shape = list(x.shape)
    for axis in self._shared_axes:
      mask_shape[axis] = 1
    keep_prob = 1.0 - rate
    keep = fastmath.random.bernoulli(rng, keep_prob, tuple(mask_shape))
    mask = keep.astype(x.dtype) / keep_prob
    return x * mask
Executes this layer as part of a forward pass through the model.

    Args:
      x: Tensor of activations.

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