Issue 1:
Description

The output of compute_pg is a vector . If I just evaluate f_x below at the desired input it works but the jacobian doesnt seem to work. I am not sure what is happening here.

@partial(jit, static_argnums=(0,))
def proj(self,p,C):
return projection_polyhedron(p,C,check_feasible = False)

@partial(jit, static_argnums=(0,))
def compute_pg(self,p,cx,cy,x_obs_inp,y_obs_inp,vx_obs_inp, vy_obs_inp):

p = jnp.reshape(p,(jnp.shape(p)[0],1))
b_obs = jnp.hstack((x_obs_inp+self.p_obs_x_ub,-(x_obs_inp-self.p_obs_x_lb),
y_obs_inp+self.p_obs_y_ub,-(y_obs_inp-self.p_obs_y_lb)))

b_obs = jnp.reshape(b_obs,(jnp.shape(self.C_obs)[0],1))
pg = ProjectedGradient(fun=self.compute_obstacle_penalty_temp,projection=self.proj,jit=True)
pg_sol = pg.run(p,hyperparams_proj=(self.A_obs,self.a_obstacle,self.C_obs,b_obs),c_x=cx,c_y=cy,vx_obs_inp=vx_obs_inp,vy_obs_inp=vy_obs_inp).params
return pg_sol

f_x = lambda cx: self.compute_pg(p[0,:],cx,best_cy.T,x_obs_inp[0,:],y_obs_inp[0,:],vx_obs_inp[0,:], vy_obs_inp[0,:])
dp_dcx = jax.jacobian(f_x)(best_cx.T)

This is the error I get:

File "/home/ims/ros2_ws/install/mpc_python/lib/python3.8/site-packages/mpc_python/mpc_expert_bilevel.py", line 592, in compute_bilevel
dp_dcx = jax.jacobian(f_x)(best_cx.T)
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.

What jax/jaxlib version are you using?

0.3.22

Which accelerator(s) are you using?

GPU

Additional system info

3.8.10,Ubuntu 20.04.4 LTS,1th Gen Intel® Core™ i9-11980HK @ 2.60GHz × 16,NVIDIA Corporation / NVIDIA GeForce RTX 3080 Laptop GPU/PCIe/SSE2

NVIDIA GPU info

-----------------------------------------------------------------------------+

| NVIDIA-SMI 510.85.02 Driver Version: 510.85.02 CUDA Version: 11.6 |

|-------------------------------+----------------------+----------------------+

| GPU Name Persistence-M| Bus-Id Disp.A | Volatile Uncorr. ECC |

| Fan Temp Perf Pwr:Usage/Cap| Memory-Usage | GPU-Util Compute M. |

| | | MIG M. |

|===============================+======================+======================|

| 0 NVIDIA GeForce ... Off | 00000000:01:00.0 Off | N/A |

| N/A 56C P5 21W / N/A | 1335MiB / 16384MiB | 5% Default |

| | | N/A |

+-------------------------------+----------------------+----------------------+

+-----------------------------------------------------------------------------+

| Processes: |

| GPU GI CI PID Type Process name GPU Memory |

| ID ID Usage |

|=============================================================================|

| 0 N/A N/A 1203 G /usr/lib/xorg/Xorg 134MiB |

| 0 N/A N/A 1845 G /usr/lib/xorg/Xorg 723MiB |

| 0 N/A N/A 2168 G /usr/bin/gnome-shell 101MiB |

| 0 N/A N/A 3875 G /usr/lib/firefox/firefox 356MiB |

| 0 N/A N/A 221513 G gnome-control-center 3MiB |

+-----------------------------------------------------------------------------+

Issue 2:
Please:

[x ] Check for duplicate issues.

[ x] Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:

# clone jkonet from github
# follow jkonet setup instructions
$ python main.py --out_dir results --config_folder configs --task semicircle

[x ] If applicable, include full error messages/tracebacks.

$ python main.py --out_dir results --config_folder configs --task semicircle

/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/chex/_src/pytypes.py:37: FutureWarning: jax.tree_structure is deprecated, and will be removed in a future release. Use jax.tree_util.tree_structure instead.
PyTreeDef = type(jax.tree_structure(None))
Started run.
/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/optax/_src/transform.py:319: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
mu = jax.tree_map( # First moment
/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/optax/_src/transform.py:321: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
nu = jax.tree_map(jnp.zeros_like, params) # Second moment
0%| | 0/1001 [00:00<?, ?it/s]/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/optax/_src/transform.py:82: FutureWarning: jax.tree_multimap is deprecated, and will be removed in a future release. Use jax.tree_util.tree_multimap instead.
return jax.tree_multimap(
/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/tree_util.py:205: FutureWarning: jax.tree_util.tree_multimap() is deprecated. Please use jax.tree_util.tree_map() instead as a drop-in replacement.
warnings.warn('jax.tree_util.tree_multimap() is deprecated. Please use jax.tree_util.tree_map() '
/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/optax/_src/transform.py:99: FutureWarning: jax.tree_multimap is deprecated, and will be removed in a future release. Use jax.tree_util.tree_multimap instead.
return jax.tree_multimap(
/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/optax/_src/transform.py:106: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
return jax.tree_map(lambda t: t / bias_correction.astype(t.dtype), moment)
/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/optax/_src/transform.py:331: FutureWarning: jax.tree_multimap is deprecated, and will be removed in a future release. Use jax.tree_util.tree_multimap instead.
updates = jax.tree_multimap(
/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/optax/_src/transform.py:359: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
updates = jax.tree_map(lambda g: step_size * g, updates)
/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/optax/_src/update.py:42: FutureWarning: jax.tree_multimap is deprecated, and will be removed in a future release. Use jax.tree_util.tree_multimap instead.
return jax.tree_multimap(
/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/optax/_src/linear_algebra.py:29: FutureWarning: jax.tree_leaves is deprecated, and will be removed in a future release. Use jax.tree_util.tree_leaves instead.
sum([jnp.sum(numerics.abs_sq(x)) for x in jax.tree_leaves(updates)]))
/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/optax/_src/clipping.py:117: FutureWarning: jax.tree_map is deprecated, and will be removed in a future release. Use jax.tree_util.tree_map instead.
updates = jax.tree_map(
/home/USER/projects/jkonet/jkonet/utils/optim.py:43: FutureWarning: jax.tree_leaves is deprecated, and will be removed in a future release. Use jax.tree_util.tree_leaves instead.
sum([jnp.sum(jnp.square(x)) for x in jax.tree_leaves(updates)]))
0%| | 0/1001 [00:00<?, ?it/s]
Traceback (most recent call last):
File "/home/USER/projects/jkonet/main.py", line 292, in <module>
main(args)
File "/home/USER/projects/jkonet/main.py", line 266, in main
run_jko(config, task_dir=task_dir, logging=args.wandb)
File "/home/USER/projects/jkonet/main.py", line 130, in run_jko
) = train_step_fn((rng_i, state_energy), batch)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/api.py", line 526, in cache_miss
out_flat = xla.xla_call(
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/core.py", line 1919, in bind
return call_bind(self, fun, *args, **params)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/core.py", line 1935, in call_bind
outs = top_trace.process_call(primitive, fun_, tracers, params)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/core.py", line 687, in process_call
return primitive.impl(f, *tracers, **params)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/dispatch.py", line 199, in _xla_call_impl
compiled_fun = xla_callable(fun, device, backend, name, donated_invars,
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/linear_util.py", line 295, in memoized_fun
ans = call(fun, *args)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/dispatch.py", line 248, in _xla_callable_uncached
return lower_xla_callable(fun, device, backend, name, donated_invars, False,
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/profiler.py", line 294, in wrapper
return func(*args, **kwargs)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/dispatch.py", line 293, in lower_xla_callable
jaxpr, out_type, consts = pe.trace_to_jaxpr_final2(
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/profiler.py", line 294, in wrapper
return func(*args, **kwargs)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/interpreters/partial_eval.py", line 2167, in trace_to_jaxpr_final2
jaxpr, out_type, consts = trace_to_subjaxpr_dynamic2(fun, main)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/interpreters/partial_eval.py", line 2117, in trace_to_subjaxpr_dynamic2
ans = fun.call_wrapped(*in_tracers_)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/linear_util.py", line 168, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/home/USER/projects/jkonet/jkonet/models/model_jko.py", line 153, in step_fn
loss_energy, loss_psi, grad_norm) = jax.lax.scan(
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/lax/control_flow/loops.py", line 259, in scan
init_flat, carry_avals, carry_avals_out, init_tree, *rest = _create_jaxpr(init)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/lax/control_flow/loops.py", line 245, in _create_jaxpr
jaxpr, consts, out_tree = _initial_style_jaxpr(
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/util.py", line 273, in wrapped
result = call(weak_arg, *args, **kwargs)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/lax/control_flow/common.py", line 57, in _initial_style_jaxpr
jaxpr, consts, out_tree = _initial_style_open_jaxpr(
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/util.py", line 273, in wrapped
result = call(weak_arg, *args, **kwargs)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/lax/control_flow/common.py", line 51, in _initial_style_open_jaxpr
jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, in_avals, debug)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/profiler.py", line 294, in wrapper
return func(*args, **kwargs)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/interpreters/partial_eval.py", line 2074, in trace_to_jaxpr_dynamic
jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/interpreters/partial_eval.py", line 2089, in trace_to_subjaxpr_dynamic
ans = fun.call_wrapped(*in_tracers_)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/linear_util.py", line 168, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/home/USER/projects/jkonet/jkonet/models/model_jko.py", line 144, in _through_time
batch = jax.lax.cond(
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/traceback_util.py", line 162, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/lax/control_flow/conditionals.py", line 254, in cond
return _cond(*args, **kwargs)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/lax/control_flow/conditionals.py", line 218, in _cond
jaxprs, consts, out_trees = _initial_style_jaxprs_with_common_consts(
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/util.py", line 220, in wrapper
return cached(config._trace_context(), *args, **kwargs)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/util.py", line 213, in cached
return f(*args, **kwargs)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/lax/control_flow/common.py", line 73, in _initial_style_jaxprs_with_common_consts
unzip3(_initial_style_open_jaxpr(fun, in_tree, in_avals, primitive_name)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/util.py", line 67, in unzip3
for x, y, z in xyzs:
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/lax/control_flow/common.py", line 73, in <genexpr>
unzip3(_initial_style_open_jaxpr(fun, in_tree, in_avals, primitive_name)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/util.py", line 273, in wrapped
result = call(weak_arg, *args, **kwargs)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/lax/control_flow/common.py", line 51, in _initial_style_open_jaxpr
jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, in_avals, debug)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/profiler.py", line 294, in wrapper
return func(*args, **kwargs)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/interpreters/partial_eval.py", line 2074, in trace_to_jaxpr_dynamic
jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/interpreters/partial_eval.py", line 2089, in trace_to_subjaxpr_dynamic
ans = fun.call_wrapped(*in_tracers_)
File "/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/linear_util.py", line 168, in call_wrapped
ans = self.f(*args, **dict(self.params, **kwargs))
File "/home/USER/projects/jkonet/jkonet/models/model_jko.py", line 146, in <lambda>
lambda x: jax.ops.index_update(x, t + 1, predicted), batch)
jax._src.traceback_util.UnfilteredStackTrace: AttributeError: module 'jax.ops' has no attribute 'index_update'

The stack trace below excludes JAX-internal frames.
The preceding is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "/home/USER/projects/jkonet/main.py", line 292, in <module>
main(args)
File "/home/USER/projects/jkonet/main.py", line 266, in main
run_jko(config, task_dir=task_dir, logging=args.wandb)
File "/home/USER/projects/jkonet/main.py", line 130, in run_jko
) = train_step_fn((rng_i, state_energy), batch)
File "/home/USER/projects/jkonet/jkonet/models/model_jko.py", line 153, in step_fn
loss_energy, loss_psi, grad_norm) = jax.lax.scan(
File "/home/USER/projects/jkonet/jkonet/models/model_jko.py", line 144, in _through_time
batch = jax.lax.cond(
File "/home/USER/projects/jkonet/jkonet/models/model_jko.py", line 146, in <lambda>
lambda x: jax.ops.index_update(x, t + 1, predicted), batch)
AttributeError: module 'jax.ops' has no attribute 'index_update'