Regularization

Dropout and stochastic depth layers.

class taktiny.nn.Dropout(p=0.5, *, broadcast_axes=(), rngs=None)

Bases: Module

Applies Dropout to the input.

Parameters:

• p (float) – The probability of an element to be zeroed. Defaults to 0.5.

• broadcast_axes (int | Sequence[int]) – Axes along which the dropout mask is broadcast. Defaults to ().

• rngs (Rngs | None) – Explicit random stream. If None, the active set_context_rng stream is looked up at call time. Evaluation and deterministic dropout cases do not require or advance a stream.

Examples

Use a temporary runtime stream without storing RNG state in the layer:

>>> from taktiny import nn
>>> import jax.numpy as jnp
>>> dropout = nn.Dropout(0.5)
>>> with nn.set_context_rng(rngs=nn.Rngs(42)):
...     y = dropout(jnp.ones((8, 4)))
>>> y.shape
(8, 4)

Or install a persistent default for the current Python context:

>>> _ = nn.set_context_rng(rngs=nn.Rngs(42))
>>> y = dropout(jnp.ones((8, 4)))
>>> _ = nn.set_context_rng(None)  # Clear the default.

Explicit rngs take priority over either form. Under JAX transformations, pass RNG state into the function, establish the scope inside it, and return the updated state; see set_context_rng for a compiled example.

class taktiny.nn.FeatureDropout(p=0.5, *, channel_axis=-1, batch_axis=0, rngs=None)

Bases: Dropout

Applies Feature Dropout (Spatial Dropout) to the input.

Parameters:

• p (float) – The probability of a feature to be zeroed. Defaults to 0.5.

• channel_axis (int) – The axis corresponding to features/channels. Defaults to -1.

• batch_axis (int | None) – The axis corresponding to the batch size, if any. Defaults to 0.

• rngs (Rngs | None) – Explicit random stream, or None to use the active set_context_rng stream at call time.

class taktiny.nn.AlphaDropout(p=0.5, *, broadcast_axes=(), rngs=None)

Bases: Module

Applies Alpha Dropout to the input, maintaining the self-normalizing property.

Parameters:

• p (float) – The probability of an element to be dropped. Defaults to 0.5.

• broadcast_axes (int | Sequence[int]) – Axes along which the dropout mask is broadcast. Defaults to ().

• rngs (Rngs | None) – Explicit random stream, or None to use the active set_context_rng stream at call time. Evaluation does not consume randomness.

class taktiny.nn.StochasticDepth(p, mode='row', *, batch_axis=0, rngs=None)

Bases: Dropout

Applies Stochastic Depth to the input.

Parameters:

• p (float) – The probability of dropping a path.

• mode (Literal['batch', 'row']) – The stochastic depth mode, either ‘batch’ or ‘row’. Defaults to ‘row’.

• batch_axis (int) – The axis corresponding to the batch size. Defaults to 0.

• rngs (Rngs | None) – Explicit random stream, or None to use the active set_context_rng stream at call time.