Provide as much information as possible. At least, this should include a description of your issue and steps to reproduce the problem. If possible also provide a summary of what steps or workarounds you have already tried.
Problem you have encountered:
If one saves & reloads a state as follows:
model_file = os.path.join(temp_dir, "flax_model")
with open(model_file, "wb") as f:
model_bytes = to_bytes(state)
f.write(model_bytes)
with open(model_file, "rb") as state_f:
state = from_bytes(AwesomeModel, state_f.read())The state previously contained weights of type jaxlib.xla_extension.DeviceArray, but is then changed to simply numpy arrays: np.ndarray.
This can lead to silent bugs when, e.g., an embedding layer flax.linen.Embed is used since it behaves differently depending on the format of its weights.
What you expected to happen:
I would expect saved weights to be loaded in the exact same format. Or maybe am I using the wrong save/load functions?
Steps to reproduce:
If you run the following code, you can see how saving/loading can lead to unexpected errors:
#!/usr/bin/env python3
import flax.linen as nn
from flax.serialization import to_bytes, from_bytes
import numpy as np
import jax.numpy as jnp
from jax.random import PRNGKey
import tempfile
import os
class AwesomeModel(nn.Module):
dtype: jnp.dtype = jnp.float32
num_embeddings = 20
embedding_size = 16
def setup(self):
self.word_emb = nn.Embed(self.num_embeddings, self.embedding_size, dtype=self.dtype)
def __call__(self, input_ids):
return self.word_emb(input_ids)
input_ids = jnp.broadcast_to(np.arange(10)[None, :], (4, 10)) # => shape [4, 10]
model = AwesomeModel()
state = model.init(PRNGKey(0), input_ids)
# No problem here!
embeddings = model.apply(state, input_ids)
# Now let's save & reload the weights
with tempfile.TemporaryDirectory() as temp_dir:
model_file = os.path.join(temp_dir, "flax_model")
with open(model_file, "wb") as f:
model_bytes = to_bytes(state)
f.write(model_bytes)
with open(model_file, "rb") as state_f:
state = from_bytes(AwesomeModel, state_f.read())
# ...and rerun the forward pass
embeddings = model.apply(state, input_ids) # Error - this is a bit unexpected
Provide as much information as possible. At least, this should include a description of your issue and steps to reproduce the problem. If possible also provide a summary of what steps or workarounds you have already tried.
Problem you have encountered:
If one saves & reloads a
stateas follows:The
statepreviously contained weights of typejaxlib.xla_extension.DeviceArray, but is then changed to simply numpy arrays:np.ndarray.This can lead to silent bugs when, e.g., an embedding layer
flax.linen.Embedis used since it behaves differently depending on the format of its weights.What you expected to happen:
I would expect saved weights to be loaded in the exact same format. Or maybe am I using the wrong save/load functions?
Steps to reproduce:
If you run the following code, you can see how saving/loading can lead to unexpected errors: