Model portability issue across machines

General
1

Hello,

Together with a collaborator, we are using BayesFlow for amortised posterior estimation and have run into a model portability issue that we have not been able to resolve. We would be very grateful for any guidance.
We have a trained BayesFlow model (BasicWorkflow with CouplingFlow, 12 layers, hidden dim 256, no summary network) that works correctly on Machine A. We are trying to load it on Machine B to run inference on new data without retraining.

Machine A: BayesFlow 2.0.3, Keras 3.10.0, TensorFlow 2.16.1, Apple M1 Max (Metal backend)
Machine B: BayesFlow 2.0.3, Keras 3.10.0, TensorFlow 2.16.2, Ubuntu 26.04 LTS (x86_64, CPU)

When we try to load the complete full_model.keras on Machine B, keras.models.load_model fails because the saved configuration contains standardize=['inference_variables'], which is passed to ContinuousApproximator, but the ContinuousApproximator implementation on Machine B does not accept this argument.

Interestingly, although both machines report BayesFlow 2.0.3, the ContinuousApproximator.__init__ signatures differ: Machine A accepts standardize, whereas Machine B does not.

We therefore tried to bypass the serialization issue by rebuilding the architecture on Machine B and loading only the inference-network weights. However, we found that the resulting architecture has 600,768 trainable parameters, compared with 600,865 on Machine A — a difference of 97 parameters. These 97 parameters appear to correspond to the standardization components (48 means + 48 variances + 1 additional parameter), which are present on Machine A but not correctly reproduced on Machine B despite explicitly setting standardize='all'.

Although the inference-network weights can be loaded, the resulting model produces incorrect posteriors. For example, the posterior for the precision parameter barely moves from the prior on Machine B, whereas on Machine A it is correctly concentrated well above the prior mean.

What we have tried

  • Saving/loading the full model with keras.models.load_model
  • Saving/loading weights only with save_weights / load_weights
  • Saving inference-network weights as NumPy .npz and loading them with set_weights
  • Explicitly setting standardize='all' on Machine B
  • Verifying that all other inputs (feature vectors and normalization statistics) are identical

All approaches either fail during loading or produce incorrect inference results.
Is there a recommended way to save and load a trained BayesFlow model portably across machines with different operating systems and hardware backends? In particular, is there a way to ensure that the standardization layers are serialized and reconstructed consistently, so that we can reproduce on Machine B exactly the trained model from Machine A?

Thank you very much for your help.

Best wishes,
Amanda Lenzi