How to retrieve the learning rate from an optimizer?

[PabloAMC]

Is there a way to look into an optimizer or opt_state to retrieve the learning rate? For example

from optax import adam
learning_rate = 1e-3

tx = adam(learning_rate = learning_rate) # outputs class `GradientTransformation`
opt_state = tx.init(params)

I have tried to look into both, but without much success.
Thanks!

[jakevdp]

Can you add imports to your example code? It's not clear which library you're importing adam from (The only one defined in jax is jax.example_libraries.optimizers.adam, but this does not accept a learning_rate argument).

[PabloAMC]

I updated the code to indicate that it is optax, thanks for highlighting!

[jakevdp]

You might ask this in the optax discussion forum (https://github.com/deepmind/optax/discussions) but looking at the source code, it doesn't look like there's any way to retrieve the parameters from a constructed optimizer.

The optimizer construction returns a chain: https://github.com/deepmind/optax/blob/84c644939d0ffe161403607a40996093e0b6c214/optax/_src/alias.py#L246-L250

And chain doesn't explicitly store any of the input parameters, but rather returns an opaque collection of python functions that close over these inputs: https://github.com/deepmind/optax/blob/84c644939d0ffe161403607a40996093e0b6c214/optax/_src/combine.py#L68