maybe_dropout

Contents

maybe_dropout#

penzai.nn.dropout.maybe_dropout(drop_rate: float | None, share_across_axes: tuple[str, ...] = ()) 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.

Returns:

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