def maybe_dropout(
    drop_rate: float | None,
    share_across_axes: tuple[str, ...] = (),
) -> StochasticDropout | DisabledDropout:
  """Constructs either a stochastic or disabled dropout layer.

  Args:
    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.
  """
  if drop_rate is None:
    return DisabledDropout(drop_rate=0.0, share_across_axes=share_across_axes)
  else:
    return StochasticDropout(
        drop_rate=drop_rate, share_across_axes=share_across_axes
    )
Constructs either a stochastic or disabled dropout layer.

  Args:
    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.
  
penzai/deprecated/v1/nn/dropout.py