Best practice of dealing with sporadic FrozenDict conversions?

Hi,

in my project, I have multiple instances of modules that look approximately like this:

```python
class TreeEncoder(nn.Module):
leaf_encoders: Any # pytree of nn.Modules

@nn.compact
def __call__(self, data: Any): # data is a pytree of jax.Array
return jnp.stack(jax.tree_leaves(jax.tree_map(lambda d, enc: enc(d), data, self.leaf_encoders)), axis=-1)
```

So essentially, this module has one encoder for every leaf of the input pytree and uses them to obtain a vector encoding of the entire tree. Usage could look like this:

```python
encoder = TreeEncoder({"img": ResNet(...), "vector": DenseNN(...)})
output = encoder({"img": jnp.zeros((480, 640, 3)), "vector": jnp.zeros(5)})
```

The beauty of JAX is that the pytree can be an arbitrary structure and not only dicts are possible. I make heavy use of this fact and sometimes just define my own dataclasses.

But here comes my problem: if I use a dictionary, flax will sometimes convert them into FrozenDicts and then the call to jax.tree_map fails with
```
Custom node type mismatch: expected type: , value: {'img': ..., 'vector': ...}
```

It took me a while to understand when these conversions happen, but I am fairly certain now that flax behaves as follows:
1. The inputs to `__call__` are never converted (data is always a dict)
2. The leaf_encoders are converted iff the TreeEncoder instance is created inside a `@nn.compact` call

How do I deal with this? I cannot call `self.leaf_encoders.unfreeze()` because it might not be a FrozenDict (and not even a dictionary). Is there some way I can disable the FrozenDict conversion in general? Or is it possible to make FrozenDicts and dicts compatible as arguments to `jax.tree_map`?

Thanks a lot in advance!

Best,
Tim