p-doom/crowd-code-dataset-0.2 huggingface dataset:
[content] salient_restore_failing_minimal_example.py (python)
import shutil

import jax.numpy as jnp
from flax.experimental import nnx
from flax.experimental.nnx import checkpoints

# This script demonstrates a potential issue in flax.nnx related to checkpointing
# and lazy parameter initialization.
#
# The hunch is that parameters like biases in Linear layers are only created if
# they are used (e.g. `use_bias=True`). When checkpointing a model created with
# `use_bias=False`, the bias parameter is not saved.
#
# Attempting to restore this checkpoint into a model instance created with
# `use_bias=True` is expected to fail, as the checkpointed state has a
# different structure (missing the 'bias' leaf in the PyTree) from the target model.

# 0. Setup a directory for the checkpoint
ckpt_dir = '/tmp/flax_nnx_bug_ckpt'
if os.path.exists(ckpt_dir):
    shutil.rmtree(ckpt_dir)
os.makedirs(ckpt_dir)


# 1. Define a simple nnx.Module
class MyModel(nnx.Module):
  """A simple model containing a Linear layer."""

  def __init__(self, use_bias: bool, *, rngs: nnx.Rngs):
    self.linear = nnx.Linear(10, 20, use_bias=use_bias, rngs=rngs)

  def __call__(self, x):
    return self.linear(x)


# 2. Instantiate the model with use_bias=False and save a checkpoint.
print('Step 1: Creating and checkpointing a model with use_bias=False.')
rngs = nnx.Rngs(0)
model_no_bias = MyModel(use_bias=False, rngs=rngs)

# Initialize the parameters by calling the model.
dummy_input = jnp.ones((1, 10))
model_no_bias(dummy_input)

# Display the state of the model. Note the absence of 'bias'.
state_no_bias, _ = nnx.split(model_no_bias)
print('State of model_no_bias:', state_no_bias)

# Save the checkpoint.
checkpoints.save_checkpoint(model_no_bias, ckpt_dir, 1)
print(f"Checkpoint saved to '{ckpt_dir}'.
")


# 3. Instantiate a new model with use_bias=True.
print('Step 2: Creating a new model with use_bias=True.')
model_with_bias = MyModel(use_bias=True, rngs=rngs)

# Initialize parameters. This will create the 'bias' parameter.
model_with_bias(dummy_input)

# Display the state. Note the presence of 'bias'.
state_with_bias, _ = nnx.split(model_with_bias)
print('State of model_with_bias (before restore):', state_with_bias)
print('')


# 4. Attempt to restore the checkpoint into the new model.
# This is expected to fail because the checkpoint is missing the 'bias' parameter
# that exists in 'model_with_bias'.
print('Step 3: Attempting to restore the checkpoint.')
try:
  checkpoints.restore_checkpoint(model_with_bias, ckpt_dir)
  print('Restore successful. (This was not expected)')
except Exception as e:
  print('Restore failed as expected.')
  print(f'Error: {e}')
finally:
  # Clean up the checkpoint directory.
  shutil.rmtree(ckpt_dir)
[content] test.ipynb (python)
import os
import shutil

import jax.numpy as jnp
from flax import nnx
import orbax.checkpoint as ocp

# The hunch is that parameters like biases in Linear layers are only created if
# they are used (e.g. `use_bias=True`). When checkpointing a model created with
# `use_bias=False`, the bias parameter is not saved.
#
# Attempting to restore this checkpoint into a model instance created with
# `use_bias=True` is expected to fail, as the checkpointed state has a
# different structure (missing the 'bias' leaf in the PyTree) from the target model.

ckpt_dir = '/tmp/flax_nnx_bug_ckpt'
if os.path.exists(ckpt_dir):
    shutil.rmtree(ckpt_dir)
os.makedirs(ckpt_dir)


class MyModel(nnx.Module):
  """A simple model containing a Linear layer."""

  def __init__(self, use_bias: bool, *, rngs: nnx.Rngs):
    self.linear = nnx.Linear(10, 20, use_bias=use_bias, rngs=rngs)

  def __call__(self, x):
    return self.linear(x)


rngs = nnx.Rngs(0)
model_no_bias = MyModel(use_bias=False, rngs=rngs)

dummy_input = jnp.ones((1, 10))
model_no_bias(dummy_input)