setu.create_mixed_nle#

setu.create_mixed_nle(x_dim, theta_dim, num_discrete, num_categories, *, condition_dim=0, flow_type='nsf', hidden_dims=(64, 64), n_flow_layers=5, n_made_layers=2, made_hidden_features=64, combiner_width=64, nn_activation=<PjitFunction of <function silu>>, log_transform_continuous=False, discrete_dropout_rate=0.2, knots=8, interval=4.0, key)[source]#

Create MixedNLE with sensible defaults.

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

Parameters:

  • x_dim (int) – Total dimension of observations (continuous + discrete).

  • theta_dim (int) – Dimension of parameters.

  • num_discrete (int) – Number of discrete variables. Must be in last columns of x.

  • num_categories (int | Sequence[int] | Array) – Number of categories per discrete variable: an int (applied to every discrete variable), a sequence of ints, or an array of shape (num_discrete,). E.g., 2 for a binary choice, [2, 3] for binary and ternary variables. A variable taking values 0..K-1 has K categories (counts 0..35 -> num_categories=36).

  • 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 flow transformer networks.

  • n_flow_layers (int) – Number of flow layers for continuous part.

  • n_made_layers (int) – Number of MADE layers for discrete part.

  • made_hidden_features (int) – Hidden features for MADE network.

  • combiner_width (int) – Width of MLP combining discrete + theta.

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

  • log_transform_continuous (bool) – Apply log-transform to positive continuous data. Useful for reaction time data.

  • discrete_dropout_rate (float) – Dropout rate for discrete net regularization.

  • 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:

MixedNLE

Note

The autoregressive mask pattern of the discrete network is deterministic (built from a fixed internal seed); key only affects network weight initialization.

Return type:

MixedNLE

Returns:

Untrained MixedNLE instance ready for fit_mixed_nle().

Raises:

  • ValueError – If x_dim <= num_discrete (no continuous variables).

  • ValueError – If num_discrete <= 0 (use NLE for pure continuous).

  • ValueError – If len(num_categories) != num_discrete.

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

Parameters:

  • x_dim (int)

  • theta_dim (int)

  • num_discrete (int)

  • num_categories (int | Sequence[int] | Array)

  • condition_dim (int)

  • flow_type (Literal['nsf', 'maf'])

  • hidden_dims (tuple[int, ...])

  • n_flow_layers (int)

  • n_made_layers (int)

  • made_hidden_features (int)

  • combiner_width (int)

  • nn_activation (Callable)

  • log_transform_continuous (bool)

  • discrete_dropout_rate (float)

  • knots (int)

  • interval (float)

  • key (Array)