Clip value of a variable in during Optimization Step?

[adam-hartshorne]

After setting up an optimization function as following. What I would like to do as indicated at the point in the code is once the gradients have been applied is to clip a particular variable values to be in a certain range. get_params(opt_state) will get you a dictionary of the variable values that I can then clip, but I don't know how (or if it is even possible) to then take the updated dictionary of variables and overwrite the opt_state.

def perform_optimisation(cost_func_from_params,
                         params,
                         learning_rate,
                         num_iterations,
                         refresh_iters=-1,
                         natural_grads_correction_func=None,
                         optimizer_maker_func=None):
    if optimizer_maker_func is None:
        opt_init, opt_update, get_params = optimizers.adam(step_size=learning_rate)
    else:
        opt_init, opt_update, get_params = optimizer_maker_func(step_size=learning_rate)

    opt_state = opt_init(params)

    def take_step(step, opt_state):
        value, grads = value_and_grad(cost_func_from_params)(get_params(opt_state))

        if jnp.any(jnp.isnan(value)):
            print('HAVE NANs :(')
            print(get_params(opt_state))
            c = cost_func_from_params(get_params(opt_state))
            raise Exception('NaN in cost function.')

        if natural_grads_correction_func is not None:
            grads = natural_grads_correction_func(grads, get_params(opt_state))

        opt_state = opt_update(step, grads, opt_state)

        ###############################################################################
        #  Want here to clip value of a variable in the opt_state to be between a certain range  #
        ############################################################################### 

        return value, opt_state

    objective_values = []

    for step in range(num_iterations):
        value, opt_state = take_step(step, opt_state)
        if refresh_iters > 0 and step % refresh_iters == 0:
            print('%5d : %0.5e' % (step, value))
        objective_values.append(value)

    return get_params(opt_state), objective_values

[cgarciae]

I think you need to do something like this:

params = get_params(opt_state)
params = jax.tree_map(lambda x: jnp.clip(x, min_val, max_val), params)
# set params back into opt_state

Which framework are you using? I think this is a bit nicer to do in optax.

[adam-hartshorne]

Thanks. I have just found optax, that seems a more flexible option