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.