optax.with_extra_args_support