I've been using jaxopt.implicit_diff.custom_root for

<a class="user-mention notranslate" data-hovercard-type="user" data-hovercard-url="/us

Yes, this looks like the same issue as <a class="issue-link js-issue-link" data-error-

CustomVJPException plus memory leak when using a for loop instead of a scan. about jaxopt HOT 3 OPEN

google commented on May 5, 2024

CustomVJPException plus memory leak when using a for loop instead of a scan.

from jaxopt.

Comments (3)

mblondel commented on May 5, 2024

@shoyer may have an idea.

from jaxopt.

shoyer commented on May 5, 2024

Yes, this looks like the same issue as #31.

You definitely need the dependency on params inside the optimality_fun passed to jaxopt.implicit_diff.custom_root. Otherwise jaxopt is not going to calculate the gradients corectly.

I think something like this should work, but it results in a different strange error:

def implicit_diff_3(params, R_init, box_size, use_for_loop=True):
    energy_fn = energy.soft_sphere_pair(displacement, **params)
    force_fn = jit(quantity.force(energy_fn))

    def optimality_fun(sol, params):
      energy_fn = energy.soft_sphere_pair(displacement, **params)
      force_fn = jit(quantity.force(energy_fn))
      return force_fn(sol)

    def solver(params, x):
        del params
        return run_minimization_scan(force_fn, x, shift, use_for_loop, num_steps = 19400)

    decorated_solver = custom_root(optimality_fun)(solver)
    R_final = decorated_solver(params, R_init)

    return (energy_fn(R_final,**params), jnp.amax(jnp.abs(force_fn(R_final,**params))))

NotImplementedError: Differentiation rule for 'custom_lin' not implemented

from jaxopt.

mblondel commented on May 5, 2024

I'm a bit confused by the variable names above. In JAXopt, we usually use optimality_fun(params, hyperparams) and solver_fun(init_params, hyperparams), where params is what is optimized and hyperparams is what is differentiated. See e.g. this example.

from jaxopt.

Recommend Projects