def __init__(self, rate=0.0, shared_axes=None, mode='train'):
    """Creates a dropout layer with the given target drop rate.

    Args:
      rate: Stochastic rate (probability) for dropping an activation value
          from the preceding layer (setting it to zero).
      shared_axes: List of axes on which the mask is shared.
      mode: If `'train'`, this layer will perform dropout; else, it will pass
          all values through unaltered.
    """
    super().__init__()
    self._initial_rate = rate
    self._shared_axes = [] if shared_axes is None else shared_axes
    self._mode = mode
Creates a dropout layer with the given target drop rate.

    Args:
      rate: Stochastic rate (probability) for dropping an activation value
          from the preceding layer (setting it to zero).
      shared_axes: List of axes on which the mask is shared.
      mode: If `'train'`, this layer will perform dropout; else, it will pass
          all values through unaltered.
    
trax/layers/core.py