Flax text-classification multi-optimizer incorrect
The below code is incorrect:

https://github.com/huggingface/transformers/blob/8d171628fe84bdf92ee40b5375d7265278180f14/examples/flax/text-classification/run_flax_glue.py#L173-L188

- The `dict` from `traverse_util.flatten_dict` has keys which are tuples of strings, not one long string with the path separated by periods.
- `optax.masked` applies the transformation wherever the mask is True, so the masks are flipped.
- Flax's LayerNorm calls the scale parameter `scale` not `weight`

I believe following code would be correct:

```py
# We use Optax's "masking" functionality to create a multi-optimizer, one
# with weight decay and the other without. Each sub-optimizer will be applied
# wherever the mask is True
decay_path = lambda p: p[-1] != "bias" and p[-2:] != ("LayerNorm", "scale")  # noqa: E731

tx = optax.chain(
    optax.masked(adamw(0.0), mask=traverse(lambda path, _: not decay_path(path))),
    optax.masked(adamw(weight_decay), mask=traverse(lambda path, _: decay_path(path))),
)
```

But, for the networks in that example, the parameters that shouldn't be decayed all have 1 dimension (biases and layernorm scales). L173-L188 can therefore be simplified to:

```py
# We use Optax's "masking" functionality to create a multi-optimizer, one
# with weight decay and the other without. Each sub-optimizer will be applied
# wherever the mask is True. The bias parameters and LayerNorm scale
# parameters should not be decayed, and these are the only parameters
# with 1 dimension.
tx = optax.chain(
    optax.masked(adamw(0.0), mask=partial(jax.tree_map, lambda p: p.ndim == 1)),
    optax.masked(adamw(weight_decay), mask=partial(jax.tree_map, lambda p: p.ndim != 1)),
)
```

Though the latter solution is simpler, perhaps the first one is better since it illustrates a more general way for users to construct a multi-optimizer. I'd be happy to open a PR with either.

cc: @marcvanzee @patrickvonplaten 
Actually `optax.adamw` can take a `mask` directly to mask weight decay (since that's pretty common):

```py
tx = optax.adamw(learning_rate=learning_rate_fn, b1=0.9, b2=0.999, eps=1e-6, weight_decay=weight_decay, 
                 mask=partial(jax.tree_map, lambda p: p.ndim != 1))
```<|||||>Hey @n2cholas,

Thanks for diving into this! 

I think a nice solution would be:

```python
tx = optax.adamw(learning_rate=learning_rate_fn, b1=0.9, b2=0.999, eps=1e-6, weight_decay=weight_decay, 
                 mask=partial(jax.tree_map, lambda p: p[-1] != "bias" and p[-2:] != ("LayerNorm", "scale")))
```

This way, it's consistent, but we also show the user clearly that `bias` and `LayerNorm` are not included in the weight decay. It would be great if you could open a PR for this :-) I can re-run the experiments afterward
<|||||>Great, will do!