Opax

* Use `jax.tree_util.*` functions to avoid warnings.
* **Full Changelog**: https://github.com/NTT123/opax/compare/v0.2.10...v0.2.11

* `ScaleBySchedule`: correct the learning rate at initialization.

**Full Changelog**: https://github.com/NTT123/opax/compare/v0.2.8...v0.2.10

Introduce strict and non-strict modes:
- Strict mode requires pytrees to be matched.
- Non-strict mode only requires tree leaves to be matched.

* Make testcases compatible with `pax3>=0.5.0`.

- Simplify the implementation by considering gradient transformations as `StateModule`.

Introduce the class method `init`. It makes things easier to understand because we explicitly say "initialize the optimizer".

Implicit:
python
optimizer = opax.adam(1e-3)(net.parameters())

Explicit:
python
optimizer = opax.adam(1e-3).init(net.parameters())