Question about implicit differentiation with optax.adam

[ramithuh]

Hi, I was trying implicit differentiation for a bi-level optimization problem I'm working on.
I was wondering whether my code's high level structure looks correct. (I didn't notice a difference with the settingimplicit_diff = True, so I was thinking maybe I'm utilizing it incorrectly)

Both the inner, and outer objective is the same mathematical equation, but the parameter we optimize is different (seq_params or tree_params).

Any pointers/feedback on my code is appreciated !

def outer_objective(tree_params, seq_params, data):
    seqs, base_tree, metadata, temp = data
    
    inner_opt_state = inner_optimizer.init_state(seq_params, tree_params, [seqs, base_tree, metadata, metadata['tLs'][0]])
    
    for i in range(0, inner_optimizer.maxiter):
       seq_params, inner_opt_state = inner_optimizer.update(seq_params, inner_opt_state, tree_params, [seqs, base_tree, metadata, metadata['tLs'][0]])

    return compute_loss_optimized(tree_params, seq_params, seqs, base_tree, metadata, temp), seq_params

@jax.jit
def inner_objective(seq_params, tree_params, data):
    seqs, base_tree, metadata, temp = data
    
    return compute_loss_optimized(tree_params, seq_params, seqs, base_tree, metadata, temp)


inner_optimizer = OptaxSolver(opt = optax.adam(0.01,  eps_root = 1e-8), fun = inner_objective, maxiter  = 5, implicit_diff = True) 
outer_optimizer = OptaxSolver(opt = optax.adam(0.1, eps_root = 1e-8), fun = outer_objective, has_aux = True)
outer_opt_state = outer_optimizer.init_state(tree_params, seq_params, [seqs, base_tree, metadata, metadata['tLs'][0]])
jitted_outer_update = jax.jit(outer_optimizer.update)

metadata['tLs'] = [0,0.01,100,5]

for epoch in 100:
    tree_params, outer_opt_state = jitted_outer_update(tree_params, outer_opt_state, outer_opt_state.aux, [seqs, base_tree, metadata, metadata['tLs'][0]])
    seq_params = outer_opt_state.aux
    
    if(_%metadata['tLs'][3]==0):
        metadata['tLs'][0] = min(metadata['tLs'][2], metadata['tLs'][0] + metadata['tLs'][1])