maybe_dropout

Contents

maybe_dropout#

penzai.nn.dropout.maybe_dropout(drop_rate: float | None, share_across_axes: tuple[str, ...] = (), random_stream_input_name: str = 'random_stream') → StochasticDropout | DisabledDropout[source]#

Constructs either a stochastic or disabled dropout layer.

Parameters:

  • drop_rate – Probability of dropping an element. If None, dropout will be disabled entirely.

  • share_across_axes – Name or names of axes to share the dropout mask over. A single dropout mask will be broadcast across these axes.

  • random_stream_input_name – Side input key for the random stream.

Returns:

A StochasticDropout layer if drop_rate is a float, or a DisabledDropout layer with drop rate 0 if drop_rate is None.

Contents