setu.create_nre#

setu.create_nre(x_dim, theta_dim, *, condition_dim=0, hidden_features=50, n_residual_blocks=2, embedding_dim=None, theta_embedding_layers=(64, 64), x_embedding_layers=(64, 64), key)[source]#

Create NRE with sensible defaults.

Creates an untrained Neural Ratio Estimator. Use fit_nre() 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.

  • hidden_features (int) – Hidden features in residual blocks.

  • n_residual_blocks (int) – Number of residual blocks in classifier.

  • embedding_dim (int | None) – Output dimension of embedding networks. None (default) means no embeddings — raw (theta, x) is concatenated and fed directly to the classifier.

  • theta_embedding_layers (tuple[int, ...]) – Hidden layer sizes for theta embedding MLP. Only used when embedding_dim is not None.

  • x_embedding_layers (tuple[int, ...]) – Hidden layer sizes for x embedding MLP. Only used when embedding_dim is not None.

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

Return type:

NRE

Returns:

Untrained NRE instance ready for fit_nre().