multistage.save

multistage.save(filename, model, **kwargs)Source

Save the model weights and configuration to a single file.

This method separates the configuration (JSON header) from the weights (binary payload). JAX arrays in the configuration are converted to metadata schemas to preserve type and shape information without writing large binary data to the text header.

Parameters:
  • filename (str) – Path to save the file (usually ending in .eqx).

  • model (eqx.Module) – Model to save.

  • kwargs (dict) – Serializable things that were given to the constructor to make this model. Excluding params to infer for inverse method. If an input was another Neural network, it is likely best to save this separately and pass it in as an additional kwarg to load.