Assertion error when trying to take jacobian of projected gradient solution using projection_polyhedron

[Basant1861]

The output of compute_pg is a (12 ,1) 2D array . This is a minimal reproducible example.

import jax.numpy as jnp
import numpy as np
from functools import partial
import jax
from jax import jit
from jaxopt import ProjectedGradient
from jaxopt.projection import projection_polyhedron

num_obs = 6
C_obs_1 = 1*jnp.identity( num_obs)
C_obs_2 = -1*jnp.identity( num_obs)
C_obs = jnp.block([
                  [ C_obs_1,0*jnp.identity( num_obs)],
                  [ C_obs_2,0*jnp.identity( num_obs)],
                  [0*jnp.identity( num_obs), C_obs_1],
                  [0*jnp.identity( num_obs), C_obs_2]
                  ])
        
A_obs = jnp.zeros((1,jnp.shape(C_obs)[1]))
a_obstacle = jnp.zeros((1,1))

def compute_obstacle_penalty_temp(p):
    cost_obs_penalty = 1.0*jnp.linalg.norm(p)**2
    return cost_obs_penalty 

def proj(p,C):
    return projection_polyhedron(p,C,check_feasible = False)

def compute_pg(p):
    
    p = jnp.reshape(p,(jnp.shape(p)[0],1))
    b_obs = jnp.ones((jnp.shape(C_obs)[0],1))
    
    pg = ProjectedGradient(fun= compute_obstacle_penalty_temp,projection= proj,jit=True)
    pg_sol = pg.run(p,hyperparams_proj=( A_obs, a_obstacle, C_obs,b_obs)).params
    return pg_sol

def compute_bilevel():
  return jax.jacobian(compute_pg)(jnp.ones((12,1)))

compute_bilevel()

This is the error I get:

File "/home/ims/ros2_ws/src/mpc_python/mpc_python/plot_test.py", line 41, in <module>
  compute_bilevel()
File "/home/ims/ros2_ws/src/mpc_python/mpc_python/plot_test.py", line 39, in compute_bilevel
  return jax.jacobian(compute_pg)(jnp.ones((12,1)))
File "/home/ims/.local/lib/python3.8/site-packages/jax/_src/api.py", line 1362, in jacfun
  jac = vmap(pullback)(_std_basis(y))
File "/home/ims/.local/lib/python3.8/site-packages/jaxopt/_src/implicit_diff.py", line 236, in solver_fun_bwd
  vjps = root_vjp(optimality_fun=optimality_fun, sol=sol,
File "/home/ims/.local/lib/python3.8/site-packages/jaxopt/_src/implicit_diff.py", line 69, in root_vjp
  u = solve(matvec, v)
File "/home/ims/.local/lib/python3.8/site-packages/jaxopt/_src/linear_solve.py", line 193, in solve_normal_cg
  Ab = rmatvec(b)  # A.T b
File "/home/ims/.local/lib/python3.8/site-packages/jaxopt/_src/linear_solve.py", line 145, in <lambda>
  return lambda y: transpose(y)[0]
AssertionError

What I observed was that if I set implicit_diff=False in the ProjectedGradient then it works but is super slow.Kindly advice.

[jakevdp]

Thanks for the question! It's difficult to understand what's going on without being able to run the code. Do you think you could edit your question to provide a minimal reproducible example? It takes a bit more work on your part, but will make it much more likely that someone here can understand and help with your issue. Thanks!

[Basant1861]

Hello @jakevdp .....sure I will try to provide a minimal reproducible example...thanks for the reply!

[Basant1861]

Hi @jakevdp ...please check I have updated the question to include a minimal reproducible example

[jakevdp]

Thanks - this looks like it might be an issue in jaxopt rather than in jax itself (though I'm not sure... I'm not entirely famiilar with jaxopt). You might try asking your question at https://github.com/google/jaxopt.