def enable(self) -> StochasticDropout:
    """Returns a stochastic, enabled version of this layer."""
    return StochasticDropout(
        drop_rate=self.drop_rate,
        share_across_axes=self.share_across_axes,
    )
Returns a stochastic, enabled version of this layer.
penzai/deprecated/v1/nn/dropout.py