Saving class args to restore model without instantiating

[ivrolan]

Usually, in all the docs and tutorials I saw, we first need to create an abstract_model by instantiating the class of tthe nnx.Module we want to restore with the same arguments as we saved the model previously. See the following snippet from Learning-JAX :

mngr_restore = ocp.CheckpointManager(ckpt_dir_ex1) # Re-open manager

# --- Create Abstract Model for Restoration ---
def create_abstract_model():
  # Use dummy RNG key/inputs for abstract creation
  return SimpleLinear(din=din, dout=dout, rngs=nnx.Rngs(params=jax.random.key(0))) # din, dout from Ex1

abstract_model = nnx.eval_shape(create_abstract_model)

# --- Split Abstract Model to get Abstract State Structure ---
graphdef_for_restore, abstract_state = nnx.split(abstract_model)

As you see, we have to "remember" the args din, dout for the SimpleLinear class. Even I find myself, cross-checking between several files, to ensure that the args used are the same.

Is there a way to save these args when saving the model, so the only information needed to restore the model is the class type SimpleLinear?

I was thinking that we could save the args separately as a metadata dict and later retrieve it, but I was wondering if there was another prefered way to do so.

[ChromeHearts]

Usually, din & dout are considered as model configs. You can save them in json or yaml format along with your checkpoint folder. These config parameters are not necessary to be saved in every checkpoint as they don't change.

[ivrolan]

Indeed, that makes sense to me, I'll implement it myself then. I was wondering if there was a way to do that directly with the API as with Keras model.save(model_path); keras.models.load_model(model_path)

[ChromeHearts]

I am not an expert of Keras nor Flax NNX, so you can consult them directly. However my suggestion is that your model params should be defined by a separate configuration file (like config.yaml) manually. You don't really need to save it programmatically again.

The typical workflow of your trainer would be:

• Use a config.yaml: Store all hyperparameters (din, dout, etc.) in this file.
• Initialize first: Your training or inference script should always start by loading the config and instantiating the model structure.
• Decide to Restore: Once you have an initialized "empty" model, you can choose to either train from scratch or overlay restored weights from a checkpoint.

Pseudo-code Example

# --- config.yaml ---
model:
  din: 128
  dout: 10
  rng_seed: 0

# --- trainer.py ---
import yaml
from flax import nnx
import orbax.checkpoint as ocp

# 1. Load your configuration
with open('config.yaml', 'r') as f:
    config = yaml.safe_load(f)

# 2. Always instantiate the model structure first using the config so you only have single source of truth
model = SimpleLinear(
    din=config['model']['din'], 
    dout=config['model']['dout'], 
    rngs=nnx.Rngs(config['model']['rng_seed'])
)

# 3. Check for existing checkpoints
mngr = ocp.CheckpointManager(ckpt_dir)

if mngr.latest_step() is not None:
    # restore as usual, and update your `model` with restored weights
    ...
else:
    print("Starting training from scratch")