bijx.noise_model

bijx.noise_model(rng, model, scale=1, *filters, noise_fn=<function normal>)[source]

Add random noise to model parameters for testing or regularization.

Applies additive noise to all parameters matching the specified filters. Useful for testing model robustness, implementing parameter regularization, or studying sensitivity to parameter perturbations.

Parameters:

  • rng (Rngs | Array) – Random number generator or key for noise generation.

  • model – Flax NNX model to add noise to.

  • scale – Scaling factor for noise magnitude.

  • *filters – Parameter filters to select which parameters to perturb.

  • noise_fn – Function for generating noise (default: normal distribution).

Returns:

New model instance with noisy parameters.

Example

>>> noisy_model = noise_model(rng, model, scale=0.1)
>>> # Adds Gaussian noise with std=0.1 to all parameters

Note

If no filters are provided, defaults to nnx.Param.