Pytree-based Optimizers

This topic is in my mind every once in a while, it has already been discussed extensively (e.g. https://github.com/deepmind/optax/issues/197#issuecomment-982548377), but I feel it needs new life because it could resolve the last remaining quirks in optax.

Optax optimizers have well defined API and contrary to neural networks they have clear ways on how to update their state, making them perfectly suitable for pytree/dataclass interfaces. Similar to what @NeilGirdhar has done [here](https://github.com/NeilGirdhar/tjax/blob/main/tjax/_src/gradient/transforms.py), one could express Pytree version of all optimizers by wrapping functional optax with the added **benefits**:

1. Optimizer can now pass through jax's function transformation boundaries, e.g. `jit`.
2. Hyper-parameters could be updated using immutable API's like `.replace()`.
3. You could get rid of the optimizer vs `opt_state` separation.
4. You can now inspect hyper-parameter updates e.g. log the learning rate under a schedule.

### Example
For this example I'l be using Flax's `PyTreeNode` but any pytree implementation is just as good.
```python
class SGD(PyTreeNode):
learning_rate: ScalarOrSchedule
momentum: Optional[float] = None
nesterov: bool = False
accumulator_dtype: Optional[Any] = field(pytree_node=False, default=None)
opt_state: Optional[OptState] = None

@property
def tx(self):
return optax.sdg(**{k: v for k, v in vars(self).items() if k != 'opt_state'})

def init(self: A, params: Params) -> A:
return self.replace(opt_state=self.tx.init(params))

def update(
self: A, updates: Updates, params: Optional[Params] = None
) -> Tuple[Updates, A]:
updates, opt_state = self.tx.update(updates, self.opt_state, params=params)
return update, self.replace(opt_state=opt_state)

# sample usage
tx = SDG(3e-4)
tx = tx.init(params)
updates, tx = tx.update(grads)
params = optax.apply_updates(params, updates)
```

### Proposal
Given that any community shim will probably not succeed, how about a `optax.pytree` namespace (naming suggestions are welcomed) where a shim could officially live and be discussed with the core team?