problem with saving a model's state

[mdcanizo]

Hello. I am following https://flax.readthedocs.io/en/latest/guides/checkpointing.html to try to save checkpoints of a model during training, which seems straight-forward, but when trying to save the state of my model I keep getting an error message which gives me no hint as to what may be going wrong. The error occurs during execution of the save command below:

checkpoint_dir = ocp.test_utils.erase_and_create_empty('/tmp/checkpoints/')
_, state = nnx.split(model)
checkpointer = ocp.PyTreeCheckpointer()
checkpointer.save(checkpoint_dir/'state', state)

and the error is

TypeError: write(): incompatible function arguments. The following argument types are supported:
1. (self: tensorstore.TensorStore, source: Union[tensorstore.TensorStore, numpy.typing.ArrayLike]) -> tensorstore.WriteFutures

Any hints as to what I might be doing wrongly would be appreciated.

[cpgaffney1]

Could you provide a full stack trace? If that's coming from TensorStore, like I think it is, it's amazing you get that far down before raising an error.

I think your usage may be incorrect. Use instead:

with ocp.StandardCheckpointer() as ckptr:
  ckptr.save(dir, state)

Also take a look at state to make sure it is a pytree of jax.Array (or np.ndarray).