Error parsing object member "base" using orbax_checkpointer.save

[jecampagne]

Dear developpers
I have some problem saving a simple model. I have setup a snippet with the crash on Collab
here.

I have tried to follow orbax/flax latest doc.

I'm using 8-devices (mimic on a 1CPU)

import os
os.environ["XLA_FLAGS"] = '--xla_force_host_platform_device_count=8'

from flax import linen as nn
from flax.training import train_state, checkpoints
from flax import struct, serialization
import orbax.checkpoint
import optax

a simple dataset fake of random uniform 2D samples, a simple model

TrainState = train_state.TrainState

class Denoiser(nn.Module):  # keep simple to get the checkpointing error
  dtype: Any = jnp.float32

  @nn.compact
  def __call__(self, z: jnp.ndarray, train: bool = True):

    x = z.reshape((args['batch_size_p'], args['z_dim']))
    x = nn.Dense(5,name=f'Model_layers_1', dtype=self.dtype)(x)
    x = nn.relu(x)
    x = nn.Dense(args['x_dim'],name=f'Model_layers_last', dtype=self.dtype)(x)
    return x

Than some model steping and state creation

@partial(jax.pmap, axis_name='num_devices')
def denoiser_step(denoiser_state: TrainState,
                  real_data: jnp.ndarray,
                  key: PRNGKey):
  r""" 1 setp gradient descent
  """
  key, key_std, key_noise = jax.random.split(key,3)
  std = jax.random.uniform(key_std, (args['batch_size_p'], 1)) * args['x_lim'] # Uniform[0,1] (nb. data in [0,1]^2)
  noise = jax.random.normal(key_noise, (args['batch_size_p'], args['z_dim'])) * std
  noisy = real_data + noise


  def loss_fn(params):

    #denoised
    output = denoiser_state.apply_fn(
        {'params': params},
        noisy)
    #skip mechanism the denoiser learn the noise
    target = noise
    denoised = noisy - output
    y = target - denoised

    # MSE
    loss = jnp.inner(y, y) / 2.0
    loss = loss.mean()

    return loss

  # Generate data with the Generator, critique it with the Discriminator.
  grad_fn = jax.value_and_grad(loss_fn)
  loss,  grads = grad_fn(denoiser_state.params)

  # Average across the devices.
  grads = jax.lax.pmean(grads, axis_name='num_devices')
  loss = jax.lax.pmean(loss, axis_name='num_devices')

  # Update the Generator through gradient descent.
  new_denoiser_state = denoiser_state.apply_gradients(
      grads=grads)

  return new_denoiser_state, loss

@partial(jax.pmap, static_broadcasted_argnums=(1, 2))
def create_state(rng, model_cls, input_shape):
  r"""Create the training state given a model class. """
  model = model_cls()

  tx = optax.sgd(learning_rate=1e-3)

  variables = model.init(rng, jnp.ones(input_shape))

  state = TrainState.create(apply_fn=model.apply, tx=tx,
      params=variables['params'])

  return state

Than some init and 1-step

from flax.training.common_utils import shard, shard_prng_key

key = jax.random.PRNGKey(seed=args['seed'])
key_denoiser, key = jax.random.split(key, 2)
key_denoiser = shard_prng_key(key_denoiser)

denoiser_state = create_state(key_denoiser, Denoiser,
    (args['batch_size_p'], *next(data_gen).shape[1:]))


key, key_denoiser = jax.random.split(key, 2)
key_denoiser = shard_prng_key(key_denoiser)

# Shard the data to possible devices.
batch_data = shard(next(data_gen))

# Take a step with the denoiser
denoiser_state, denoiser_loss = denoiser_step(denoiser_state,
          batch_data, key_denoiser)

Then, a checkpoint

from flax.training import orbax_utils
from tqdm.notebook import tqdm
import numpy as np

ckpt = {'model': denoiser_state}
orbax_checkpointer = orbax.checkpoint.PyTreeCheckpointer()
save_args = orbax_utils.save_args_from_target(ckpt)
orbax_checkpointer.save('/denoiser/flax_ckpt/orbax/single_save', ckpt, save_args=save_args)

But this leads to a crash

WARNING:root:Conversion for <class 'jaxlib.xla_extension.PmapSharding'> has not been implemented.
WARNING:root:Conversion for <class 'jaxlib.xla_extension.PmapSharding'> has not been implemented.
WARNING:root:Conversion for <class 'jaxlib.xla_extension.PmapSharding'> has not been implemented.
WARNING:root:Conversion for <class 'jaxlib.xla_extension.PmapSharding'> has not been implemented.
WARNING:root:Conversion for <class 'jaxlib.xla_extension.PmapSharding'> has not been implemented.

ValueError: Error parsing object member "base": Error parsing object member "metadata": Error parsing object member "chunks": Array has length 0 but should have length 1 [source locations='tensorstore/internal/json_binding/std_array.h:60\ntensorstore/internal/json_binding/json_binding.h:865\ntensorstore/internal/json_binding/json_binding.h:829\ntensorstore/internal/json_binding/json_binding.h:865\ntensorstore/internal/json_binding/json_binding.h:829\ntensorstore/internal/json_binding/json_binding.h:865\ntensorstore/internal/json_binding/json_binding.h:829\ntensorstore/internal/json_binding/json_binding.h:829']

Any help is welcome. Thanks