onnx

lanfactory.onnx

Modules:

• bayesflow –

  Export trained bayesflow approximators to ONNX for HSSM consumption.

• contract –

  The single-trial ONNX contract, as a check instead of a paragraph.

• jax_export –

  Export jaxtrain networks to ONNX. Can be run as a script.

• sbi –

  Export trained sbi estimators to ONNX for HSSM consumption.

• transform_onnx –

  This module contains the function to transform Torch/Jax models to ONNX format.

Functions:

• assert_single_trial_contract –

  Raise AssertionError unless the artifact satisfies the contract.

• transform_bayesflow_to_onnx –

  Export a trained bayesflow approximator to a single-trial ONNX graph.

• transform_jax_to_onnx –

  Transform a trained JaxMLP to ONNX format.

• transform_sbi_to_onnx –

  Export a trained sbi estimator to a single-trial ONNX graph.

• transform_to_onnx –

  Transforms a TorchMLP model to ONNX format.

lanfactory.onnx.assert_single_trial_contract

assert_single_trial_contract(onnx_path: str | Path, expected_input_width: int | None = None, allowed_ops: set[str] | None = None) -> dict

Raise AssertionError unless the artifact satisfies the contract.

Parameters:

• onnx_path (str | Path) –

  The exported artifact.

• expected_input_width (int | None, default: None ) –

  The per-trial input width (the last dimension), when the caller knows it. Catches an exporter that silently changed its input layout.

• allowed_ops (set[str] | None, default: None ) –

  When given, the graph's op types must be a subset. Useful to pin a lowering that a pinned 0.x exporter dependency could change under you.

Returns:

• dict –

  {"input_shape", "input_width", "ops"} for further assertions.

lanfactory.onnx.transform_bayesflow_to_onnx

transform_bayesflow_to_onnx(approximator: Any, path: str, *, mode: Literal['nle', 'nre'] = 'nle', example_theta_dim: int, example_x_dim: int, opset: int = 17) -> None

Export a trained bayesflow approximator to a single-trial ONNX graph.

Parameters:

• approximator (Any) –

  Trained bayesflow approximator. For mode="nle" this is a :class:bayesflow.ContinuousApproximator whose inference_variables slot was trained on the observation x (with inference_conditions holding the parameters θ). For mode="nre" this is a :class:bayesflow.RatioApproximator trained with the opposite convention (inference_variables=θ, inference_conditions=x).

• path (str) –

  Filesystem path to write the .onnx artifact to.

• mode (Literal['nle', 'nre'], default: 'nle' ) –

  "nle" exports log p(x|θ) with the standardizer Jacobian baked in. "nre" exports the classifier logit as the log-likelihood up to a θ-independent constant (which drops out in MCMC).

• example_theta_dim (int) –

  Parameter-vector dimensionality used to trace the graph.

• example_x_dim (int) –

  Observation-vector dimensionality used to trace the graph.

• opset (int, default: 17 ) –

  ONNX opset version. Pinned to 17 by default for reproducibility against jaxonnxruntime.

Raises:

• RuntimeError –

  If KERAS_BACKEND is not "torch".

• TypeError –

  If mode is "nle" and the approximator does not have an inference_network with .log_prob; or if mode is "nre" and the approximator does not have a projector.

• ValueError –

  If the approximator's adapter contains any non-trivial transforms, if example_theta_dim or example_x_dim is not positive, or if mode is not one of "nle"/"nre".

lanfactory.onnx.transform_jax_to_onnx

transform_jax_to_onnx(network_config_file: str, state_file: str, input_shape: int, output_onnx_file: str, opset: int = DEFAULT_OPSET) -> None

Transform a trained JaxMLP to ONNX format.

Arguments

network_config_file (str):
    Path to the pickle file containing the network configuration
    (``layer_sizes``, ``activations``, ``train_output_type``).
state_file (str):
    Path to the ``*_train_state.jax`` file written by the jax trainer
    (flax ``to_bytes`` serialization of the parameters).
input_shape (int):
    The size of the single-trial input vector for the model
    (``n_params + 2`` for LANs).
output_onnx_file (str):
    Path to the output ONNX file.
opset (int):
    ONNX opset version to target.

lanfactory.onnx.transform_sbi_to_onnx

transform_sbi_to_onnx(estimator: Module, path: str, *, mode: Literal['nle', 'nre'] = 'nle', example_theta_dim: int, example_x_dim: int, opset: int = 17) -> None

Export a trained sbi estimator to a single-trial ONNX graph.

Parameters:

• estimator (Module) –

  A trained sbi estimator. For mode="nle" this is a ConditionalDensityEstimator (as returned by NLE_A.train()); for mode="nre" it is a ratio-estimator classifier (from NRE_A/B/ C, BNRE).

• path (str) –

  Filesystem path to write the .onnx artifact to.

• mode (Literal['nle', 'nre'], default: 'nle' ) –

  "nle" exports estimator.log_prob as the log-likelihood with the standardization Jacobian baked in. "nre" exports the classifier logit as the log-likelihood up to a θ-independent constant (lands in C4).

• example_theta_dim (int) –

  Parameter-vector dimensionality used to trace the graph.

• example_x_dim (int) –

  Observation-vector dimensionality used to trace the graph.

• opset (int, default: 17 ) –

  ONNX opset version. Pinned to 17 by default for reproducibility against jaxonnxruntime.

Notes

Only likelihood-shaped families are supported. NPE/posterior estimators are rejected by convention (the caller asserts mode="nle" only for true likelihood estimators). Score-based / flow-matching estimators (FMPE, NPSE), TabPFN-based estimators, and neural spline flows (blocked on missing SearchSorted in jaxonnxruntime) are rejected with a clear error.

lanfactory.onnx.transform_to_onnx

transform_to_onnx(network_config_file: str, state_dict_file: str, input_shape: int, output_onnx_file: str) -> None

Transforms a TorchMLP model to ONNX format.

Arguments

network_config_file (str):
    Path to the pickle file containing the network configuration.
state_dict_file (str):
    Path to the file containing the state dictionary of the model.
input_shape (int):
    The size of the input tensor for the model.
output_onnx_file (str):
    Path to the output ONNX file.