setu.create_nle#

setu.create_nle(x_dim, theta_dim, *, condition_dim=0, flow_type='nsf', hidden_dims=(64, 64), n_layers=5, nn_activation=<PjitFunction of <function silu>>, knots=8, interval=4.0, key)[source]#

Create NLE with sensible defaults.

Creates an untrained Neural Likelihood Estimator. Use fit_nle() to train.

Parameters:

  • x_dim (int) – Dimension of observations.

  • theta_dim (int) – Dimension of parameters.

  • condition_dim (int) – Dimension of experimental conditions. Use 0 (default) if no conditions.

  • flow_type (Literal['nsf', 'maf']) – Flow architecture. “nsf” for Neural Spline Flow (recommended), “maf” for Masked Autoregressive Flow.

  • hidden_dims (tuple[int, ...]) – Hidden layer sizes for transformer networks. Determines network capacity.

  • n_layers (int) – Number of flow layers. More layers = more expressive.

  • nn_activation (Callable) – Activation function for transformer networks. Default silu.

  • knots (int) – Number of spline knots for NSF. Ignored for MAF.

  • interval (float) – Spline interval for NSF. Ignored for MAF.

  • key (Array) – JAX random key for initialization.

Return type:

NLE

Returns:

Untrained NLE instance ready for fit_nle().

Raises:

ValueError – If flow_type is not “nsf” or “maf”.