Numpyro Extras

register_parameters

register_parameters(model, priors=None, prefix='')

Register GPJax parameters with Numpyro.

Parameters:

• model (Module) –

  The GPJax model (flax.nnx.Module).

• priors (Dict[str, Distribution] | None, default: None ) –

  Optional dictionary mapping parameter names to Numpyro distributions.

• prefix (str, default: '' ) –

  Optional prefix for parameter names.

Returns:

• Module –

  The model with parameters updated from Numpyro samples.