Markus Kunesch

When wrapping `optax.adafactor` with `optax.inject_hyperparams` without specifying `static_args` ``` optax.inject_hyperparams(optax.adafactor)(learning_rate=0.1) ``` the `init` function of the resulting `GradientTransformation` cannot be jit compiled. The reason is that by default `inject_hyperparams` treats...