`serialization.from_state_dict` does not restore to jax.Arrays

Deserialising a single jax.Array with `flax.serialization` does not reconstruct `jax.Array`correctly even when they are present in the target pytree.

See

```python
import jax
import jax.numpy as jnp
import flax
from flax import serialization

pt = {'a':jnp.ones((3,4))}

bdata = serialization.to_bytes(pt)
pt_loaded = serialization.from_bytes(pt, bdata)

print("Standard Deserialization output type:", type(pt_loaded['a']))
> Standard Deserialization output type:
```

This is because there is no rule for how to deal with jax.Arrays.

To fix it, one needs to register this rule
```python
serialization.register_serialization_state(
type(pt['a']),
lambda x: x,
lambda x, sd: jax.numpy.asarray(sd, dtype=x.dtype),
override=True
)
bdata = serialization.to_bytes(pt)
pt_loaded = serialization.from_bytes(pt, bdata)

print("type Deserialization output type:", type(pt_loaded['a']))
> type Deserialization output type:
```

Would it be possible to get this inside of flax itself?