def maybe_dropout(
    drop_rate: float | None,
    share_across_axes: tuple[str, ...] = (),
    random_stream_input_name: str = "random_stream",
) -> 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.
    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.
  """
  if drop_rate is None:
    return DisabledDropout(
        drop_rate=0.0,
        share_across_axes=share_across_axes,
        random_stream_input_name=random_stream_input_name,
    )
  else:
    return StochasticDropout(
        drop_rate=drop_rate,
        share_across_axes=share_across_axes,
        random_stream_input_name=random_stream_input_name,
    )
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.
    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.
  
penzai/nn/dropout.py