unravel_pytree not functioning on integer input

[dulacarnold]

Hi Jax-team,

There is a bug related to vjp of the flatten operation and integer input:

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)

returns

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

This is only if I provide integer input. A proposed fix by @mattjj would be to recast to float and raise a warning, we agreed I would file a bug here.

All the best!

Gabe