How to get params of optimizer?

[arihantbansal]

Is there any way to get parameters of the optimizer/network at a step given the optimizer state? Similar to get_params(opt_state) which exists in jax.example_libraries.optimizers

[pharringtonp19]

@arihantbansal

I think you are looking for the following two lines of code.

updates, opt_state = optimizer.update(grads, opt_state)
params = optax.apply_updates(params, updates)

At each iteration

1. Update the optimizer state optimizer.update(grads, opt_state)
2. Update the parameters of the neural network optax.apply_updates(params, updates)

[louiskenyon98]

How do you get the initial optimizer state, given parameters, for the network? When initialising with the previous method in Jax you would call get_params() to obtain the starting state of the network which is then passed to the step function.

I found my own answer: you call optimizer.init(params)