Regular dict representation not indented

After the [dict migration](https://github.com/google/flax/pull/3193), Flax now returns regular dicts when calling the `.init`, `.init_with_output` and `.apply` Module methods. However the representation of regular dicts are not as readable compared to the indented version of FrozenDicts.

Regular dicts:
```
class MLP(nn.Module):
@nn.compact
def __call__(self, x):
x = nn.Dense(5)(x)
x = nn.relu(x)
return x
x = jnp.ones((1,3))
model = MLP()
params = model.init(jax.random.PRNGKey(0), x)['params']
state = TrainState.create(apply_fn=model.apply, params=params, tx=optax.adam(1e-3))
state
```
```
TrainState(step=0, apply_fn=, params={'Dense_0': {'kernel': Array([[ 0.37229332, -0.4265755 , -1.1151816 , -0.09558704, -0.62169886],
[-1.060781 , 1.0546707 , 0.33051118, -0.7090655 , 0.37682843],
[-0.30747807, -0.39064118, -0.25515485, 0.5127583 , -0.5559202 ]], dtype=float32), 'bias': Array([0., 0., 0., 0., 0.], dtype=float32)}}, tx=GradientTransformationExtraArgs(init=.init_fn at 0x7fe6daeca4d0>, update=.update_fn at 0x7fe6daeca9e0>), opt_state=(ScaleByAdamState(count=Array(0, dtype=int32), mu={'Dense_0': {'bias': Array([0., 0., 0., 0., 0.], dtype=float32), 'kernel': Array([[0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0.]], dtype=float32)}}, nu={'Dense_0': {'bias': Array([0., 0., 0., 0., 0.], dtype=float32), 'kernel': Array([[0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0.]], dtype=float32)}}), EmptyState()))
```

FrozenDicts:
```
state = TrainState.create(apply_fn=model.apply, params=flax.core.freeze(params), tx=optax.adam(1e-3))
state
```
```
TrainState(step=0, apply_fn=, params=FrozenDict({
Dense_0: {
kernel: Array([[ 0.37229332, -0.4265755 , -1.1151816 , -0.09558704, -0.62169886],
[-1.060781 , 1.0546707 , 0.33051118, -0.7090655 , 0.37682843],
[-0.30747807, -0.39064118, -0.25515485, 0.5127583 , -0.5559202 ]], dtype=float32),
bias: Array([0., 0., 0., 0., 0.], dtype=float32),
},
}), tx=GradientTransformationExtraArgs(init=.init_fn at 0x7fe710e21ea0>, update=.update_fn at 0x7fe710e225f0>), opt_state=(ScaleByAdamState(count=Array(0, dtype=int32), mu=FrozenDict({
Dense_0: {
bias: Array([0., 0., 0., 0., 0.], dtype=float32),
kernel: Array([[0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0.]], dtype=float32),
},
}), nu=FrozenDict({
Dense_0: {
bias: Array([0., 0., 0., 0., 0.], dtype=float32),
kernel: Array([[0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0.],
[0., 0., 0., 0., 0.]], dtype=float32),
},
})), EmptyState()))
```

The indented representation can be viewed by calling [`flax.core.pretty_repr`](https://flax.readthedocs.io/en/latest/api_reference/flax.core.frozen_dict.html#flax.core.frozen_dict.pretty_repr) on the dict. Alternatively we could subclass dict and override the `__repr__` method to return an indented representation and have Flax return this subclass when `.init`, `.init_with_output` and `.apply` are called:
```
@flax.struct.dataclass
class MutableDict(dict):
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
def __repr__(self):
return 'MutableDict' + flax.core.pretty_repr(self)

state = TrainState.create(apply_fn=model.apply, params=MutableDict(params), tx=optax.adam(1e-3))
state
```
```
TrainState(step=0, apply_fn=, params=MutableDict{
Dense_0: {
kernel: Array([[ 0.37229332, -0.4265755 , -1.1151816 , -0.09558704, -0.62169886],
[-1.060781 , 1.0546707 , 0.33051118, -0.7090655 , 0.37682843],
[-0.30747807, -0.39064118, -0.25515485, 0.5127583 , -0.5559202 ]], dtype=float32),
bias: Array([0., 0., 0., 0., 0.], dtype=float32),
},
}, tx=GradientTransformationExtraArgs(init=.init_fn at 0x7fe710e3a5f0>, update=.update_fn at 0x7fe710e39990>), opt_state=(ScaleByAdamState(count=Array(0, dtype=int32), mu=MutableDict{}, nu=MutableDict{}), EmptyState()))
```
Another option is to add a section in the [dict migration guide](https://flax.readthedocs.io/en/latest/guides/regular_dict_upgrade_guide.html) to let users know they can get the indented representation by calling `flax.core.pretty_repr` (although this currently works only on FrozenDicts and regular dicts, and not other objects like TrainState).