How to use ravel_pytree in Flax/JAX?

[marcvanzee]

Original question by @dulacarnold.

I am trying a simple ravel, unravel and I'm losing the contents

pt = {'foo': jnp.array([1,2,3]), 'bar': {'baz': jnp.array([5,6]), 'bork': jnp.array([7,8])}}
flat, unflatten = jax.flatten_util.ravel_pytree(pt)
unflatten(flat)

and I get...

{'bar': {'baz': array([(b'',), (b'',)], dtype=[('float0', 'V')]),
  'bork': array([(b'',), (b'',)], dtype=[('float0', 'V')])},
 'foo': array([(b'',), (b'',), (b'',)], dtype=[('float0', 'V')])}

[marcvanzee]

You have to use this function with floats, with integers it will have this problem.