jax.Array must be fully replicated to be saved in aggregate file

[tatami-galaxy]

I'm trying to save a checkpoint in Flax and getting this error message. Saving code :

ckpt = {'state': state, 'config': model.config} save_args = orbax_utils.save_args_from_target(ckpt) checkpoint_manager.save(global_step + 1, ckpt, save_kwargs={'save_args': save_args})

What could be causing this? I tried disabling jax.Array with jax.config.update('jax_array', False) but that does not work with jax and jaxlib 0.4.7.

Edit : state is an instance of flax.training.train_state

[yashk2810]

Looks like this belongs in the flax repo? Can you file an issue here: https://github.com/google/flax?

[tatami-galaxy]

Yeah sorry I realised that later. I've filed an issue on the flax repo. Thanks.