<h3 dir="auto">Description</h3>
<p dir="auto">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.</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="@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)
"><pre class="notranslate"><code class="notranslate">@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)

</code></pre></div>
<p dir="auto">This is the error I get:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="File &quot;/home/ims/ros2_ws/install/mpc_python/lib/python3.8/site-packages/mpc_python/mpc_expert_bilevel.py&quot;, line 592, in compute_bilevel
  dp_dcx = jax.jacobian(f_x)(best_cx.T)
File &quot;/home/ims/.local/lib/python3.8/site-packages/jaxopt/_src/implicit_diff.py&quot;, line 236, in solver_fun_bwd
  vjps = root_vjp(optimality_fun=optimality_fun, sol=sol,
File &quot;/home/ims/.local/lib/python3.8/site-packages/jaxopt/_src/implicit_diff.py&quot;, line 69, in root_vjp
  u = solve(matvec, v)
File &quot;/home/ims/.local/lib/python3.8/site-packages/jaxopt/_src/linear_solve.py&quot;, line 193, in solve_normal_cg
  Ab = rmatvec(b)  # A.T b
File &quot;/home/ims/.local/lib/python3.8/site-packages/jaxopt/_src/linear_solve.py&quot;, line 145, in &lt;lambda&gt;
  return lambda y: transpose(y)[0]
AssertionError"><pre class="notranslate"><code class="notranslate">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 &lt;lambda&gt;
  return lambda y: transpose(y)[0]
AssertionError
</code></pre></div>
<p dir="auto">What I observed was that if I set implicit_diff=False in the ProjectedGradient then it works but is super slow.Kindly advice.</p>
<h3 dir="auto">What jax/jaxlib version are you using?</h3>
<p dir="auto">0.3.22</p>
<h3 dir="auto">Which accelerator(s) are you using?</h3>
<p dir="auto">GPU</p>
<h3 dir="auto">Additional system info</h3>
<p dir="auto">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</p>
<h3 dir="auto">NVIDIA GPU info</h3>
<p dir="auto">-----------------------------------------------------------------------------+<br>
| NVIDIA-SMI 510.85.02    Driver Version: 510.85.02    CUDA Version: 11.6     |<br>
|-------------------------------+----------------------+----------------------+<br>
| GPU  Name        Persistence-M| Bus-Id        Disp.A | Volatile Uncorr. ECC |<br>
| Fan  Temp  Perf  Pwr:Usage/Cap|         Memory-Usage | GPU-Util  Compute M. |<br>
|                               |                      |               MIG M. |<br>
|===============================+======================+======================|<br>
|   0  NVIDIA GeForce ...  Off  | 00000000:01:00.0 Off |                  N/A |<br>
| N/A   56C    P5    21W /  N/A |   1335MiB / 16384MiB |      5%      Default |<br>
|                               |                      |                  N/A |<br>
+-------------------------------+----------------------+----------------------+</p>
<p dir="auto">+-----------------------------------------------------------------------------+<br>
| Processes:                                                                  |<br>
|  GPU   GI   CI        PID   Type   Process name                  GPU Memory |<br>
|        ID   ID                                                   Usage      |<br>
|=============================================================================|<br>
|    0   N/A  N/A      1203      G   /usr/lib/xorg/Xorg                134MiB |<br>
|    0   N/A  N/A      1845      G   /usr/lib/xorg/Xorg                723MiB |<br>
|    0   N/A  N/A      2168      G   /usr/bin/gnome-shell              101MiB |<br>
|    0   N/A  N/A      3875      G   /usr/lib/firefox/firefox          356MiB |<br>
|    0   N/A  N/A    221513      G   gnome-control-center                3MiB |<br>
+-----------------------------------------------------------------------------+</p>
<p dir="auto">Please:</p>
<ul dir="auto">
<li>[x ] Check for duplicate issues.</li>
<li>[ x] Provide a complete example of how to reproduce the bug, wrapped in triple backticks like this:</li>
</ul>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="# clone jkonet from github
# follow jkonet setup instructions
$ python main.py --out_dir results --config_folder configs --task semicircle"><pre class="notranslate"><span class="pl-c"># clone jkonet from github</span>
<span class="pl-c"># follow jkonet setup instructions</span>
$ <span class="pl-s1">python</span> <span class="pl-s1">main</span>.<span class="pl-s1">py</span> <span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-s1">out_dir</span> <span class="pl-s1">results</span> <span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-s1">config_folder</span> <span class="pl-s1">configs</span> <span class="pl-c1">-</span><span class="pl-c1">-</span><span class="pl-s1">task</span> <span class="pl-s1">semicircle</span></pre></div>
<ul dir="auto">
<li>[x ] If applicable, include full error messages/tracebacks.</li>
</ul>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="$ 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&lt;?, ?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&lt;?, ?it/s]
Traceback (most recent call last):
  File &quot;/home/USER/projects/jkonet/main.py&quot;, line 292, in &lt;module&gt;
    main(args)
  File &quot;/home/USER/projects/jkonet/main.py&quot;, line 266, in main
    run_jko(config, task_dir=task_dir, logging=args.wandb)
  File &quot;/home/USER/projects/jkonet/main.py&quot;, line 130, in run_jko
    ) = train_step_fn((rng_i, state_energy), batch)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/traceback_util.py&quot;, line 162, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/api.py&quot;, line 526, in cache_miss
    out_flat = xla.xla_call(
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/core.py&quot;, line 1919, in bind
    return call_bind(self, fun, *args, **params)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/core.py&quot;, line 1935, in call_bind
    outs = top_trace.process_call(primitive, fun_, tracers, params)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/core.py&quot;, line 687, in process_call
    return primitive.impl(f, *tracers, **params)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/dispatch.py&quot;, line 199, in _xla_call_impl
    compiled_fun = xla_callable(fun, device, backend, name, donated_invars,
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/linear_util.py&quot;, line 295, in memoized_fun
    ans = call(fun, *args)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/dispatch.py&quot;, line 248, in _xla_callable_uncached
    return lower_xla_callable(fun, device, backend, name, donated_invars, False,
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/profiler.py&quot;, line 294, in wrapper
    return func(*args, **kwargs)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/dispatch.py&quot;, line 293, in lower_xla_callable
    jaxpr, out_type, consts = pe.trace_to_jaxpr_final2(
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/profiler.py&quot;, line 294, in wrapper
    return func(*args, **kwargs)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/interpreters/partial_eval.py&quot;, line 2167, in trace_to_jaxpr_final2
    jaxpr, out_type, consts = trace_to_subjaxpr_dynamic2(fun, main)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/interpreters/partial_eval.py&quot;, line 2117, in trace_to_subjaxpr_dynamic2
    ans = fun.call_wrapped(*in_tracers_)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/linear_util.py&quot;, line 168, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File &quot;/home/USER/projects/jkonet/jkonet/models/model_jko.py&quot;, line 153, in step_fn
    loss_energy, loss_psi, grad_norm) = jax.lax.scan(
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/traceback_util.py&quot;, line 162, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/lax/control_flow/loops.py&quot;, line 259, in scan
    init_flat, carry_avals, carry_avals_out, init_tree, *rest = _create_jaxpr(init)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/lax/control_flow/loops.py&quot;, line 245, in _create_jaxpr
    jaxpr, consts, out_tree = _initial_style_jaxpr(
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/util.py&quot;, line 273, in wrapped
    result = call(weak_arg, *args, **kwargs)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/lax/control_flow/common.py&quot;, line 57, in _initial_style_jaxpr
    jaxpr, consts, out_tree = _initial_style_open_jaxpr(
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/util.py&quot;, line 273, in wrapped
    result = call(weak_arg, *args, **kwargs)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/lax/control_flow/common.py&quot;, line 51, in _initial_style_open_jaxpr
    jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, in_avals, debug)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/profiler.py&quot;, line 294, in wrapper
    return func(*args, **kwargs)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/interpreters/partial_eval.py&quot;, line 2074, in trace_to_jaxpr_dynamic
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/interpreters/partial_eval.py&quot;, line 2089, in trace_to_subjaxpr_dynamic
    ans = fun.call_wrapped(*in_tracers_)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/linear_util.py&quot;, line 168, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File &quot;/home/USER/projects/jkonet/jkonet/models/model_jko.py&quot;, line 144, in _through_time
    batch = jax.lax.cond(
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/traceback_util.py&quot;, line 162, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/lax/control_flow/conditionals.py&quot;, line 254, in cond
    return _cond(*args, **kwargs)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/lax/control_flow/conditionals.py&quot;, line 218, in _cond
    jaxprs, consts, out_trees = _initial_style_jaxprs_with_common_consts(
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/util.py&quot;, line 220, in wrapper
    return cached(config._trace_context(), *args, **kwargs)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/util.py&quot;, line 213, in cached
    return f(*args, **kwargs)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/lax/control_flow/common.py&quot;, line 73, in _initial_style_jaxprs_with_common_consts
    unzip3(_initial_style_open_jaxpr(fun, in_tree, in_avals, primitive_name)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/util.py&quot;, line 67, in unzip3
    for x, y, z in xyzs:
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/lax/control_flow/common.py&quot;, line 73, in &lt;genexpr&gt;
    unzip3(_initial_style_open_jaxpr(fun, in_tree, in_avals, primitive_name)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/util.py&quot;, line 273, in wrapped
    result = call(weak_arg, *args, **kwargs)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/lax/control_flow/common.py&quot;, line 51, in _initial_style_open_jaxpr
    jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, in_avals, debug)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/_src/profiler.py&quot;, line 294, in wrapper
    return func(*args, **kwargs)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/interpreters/partial_eval.py&quot;, line 2074, in trace_to_jaxpr_dynamic
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/interpreters/partial_eval.py&quot;, line 2089, in trace_to_subjaxpr_dynamic
    ans = fun.call_wrapped(*in_tracers_)
  File &quot;/home/USER/anaconda3/envs/jko/lib/python3.9/site-packages/jax/linear_util.py&quot;, line 168, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File &quot;/home/USER/projects/jkonet/jkonet/models/model_jko.py&quot;, line 146, in &lt;lambda&gt;
    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 &quot;/home/USER/projects/jkonet/main.py&quot;, line 292, in &lt;module&gt;
    main(args)
  File &quot;/home/USER/projects/jkonet/main.py&quot;, line 266, in main
    run_jko(config, task_dir=task_dir, logging=args.wandb)
  File &quot;/home/USER/projects/jkonet/main.py&quot;, line 130, in run_jko
    ) = train_step_fn((rng_i, state_energy), batch)
  File &quot;/home/USER/projects/jkonet/jkonet/models/model_jko.py&quot;, line 153, in step_fn
    loss_energy, loss_psi, grad_norm) = jax.lax.scan(
  File &quot;/home/USER/projects/jkonet/jkonet/models/model_jko.py&quot;, line 144, in _through_time
    batch = jax.lax.cond(
  File &quot;/home/USER/projects/jkonet/jkonet/models/model_jko.py&quot;, line 146, in &lt;lambda&gt;
    lambda x: jax.ops.index_update(x, t + 1, predicted), batch)
AttributeError: module 'jax.ops' has no attribute 'index_update'
"><pre class="notranslate"><code class="notranslate">$ 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&lt;?, ?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&lt;?, ?it/s]
Traceback (most recent call last):
  File "/home/USER/projects/jkonet/main.py", line 292, in &lt;module&gt;
    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 &lt;genexpr&gt;
    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 &lt;lambda&gt;
    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 &lt;module&gt;
    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 &lt;lambda&gt;
    lambda x: jax.ops.index_update(x, t + 1, predicted), batch)
AttributeError: module 'jax.ops' has no attribute 'index_update'

</code></pre></div>