Hi! I am trying to run the finite strain FeFp plasticity tutorial, but I am getting some errors while running the script. I have installed FEniCSx (v.0.8) and Jax(v.0.5.2) using conda-forge.
JaxStackTraceBeforeTransformation Traceback (most recent call last)
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/runpy.py:198, in _run_module_as_main()
197 sys.argv[0] = mod_spec.origin
--> 198 return _run_code(code, main_globals, None,
199 "__main__", mod_spec)
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/runpy.py:88, in _run_code()
81 run_globals.update(__name__ = mod_name,
82 __file__ = fname,
83 __cached__ = cached,
(...)
86 __package__ = pkg_name,
87 __spec__ = mod_spec)
---> 88 exec(code, run_globals)
89 return run_globals
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/site-packages/ipykernel/kernelapp.py:739, in start()
738 try:
--> 739 self.io_loop.start()
740 except KeyboardInterrupt:
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/site-packages/tornado/platform/asyncio.py:205, in start()
204 def start(self) -> None:
--> 205 self.asyncio_loop.run_forever()
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/asyncio/base_events.py:683, in run_forever()
682 while True:
--> 683 self._run_once()
684 if self._stopping:
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/asyncio/base_events.py:2040, in _run_once()
2039 else:
-> 2040 handle._run()
2041 handle = None
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/asyncio/events.py:89, in _run()
88 try:
---> 89 self._context.run(self._callback, *self._args)
90 except (SystemExit, KeyboardInterrupt):
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/site-packages/ipykernel/kernelbase.py:545, in dispatch_queue()
544 try:
--> 545 await self.process_one()
546 except Exception:
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/site-packages/ipykernel/kernelbase.py:534, in process_one()
533 return
--> 534 await dispatch(*args)
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/site-packages/ipykernel/kernelbase.py:437, in dispatch_shell()
436 if inspect.isawaitable(result):
--> 437 await result
438 except Exception:
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/site-packages/ipykernel/ipkernel.py:362, in execute_request()
361 self._associate_new_top_level_threads_with(parent_header)
--> 362 await super().execute_request(stream, ident, parent)
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/site-packages/ipykernel/kernelbase.py:778, in execute_request()
777 if inspect.isawaitable(reply_content):
--> 778 reply_content = await reply_content
780 # Flush output before sending the reply.
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/site-packages/ipykernel/ipkernel.py:449, in do_execute()
448 if accepts_params["cell_id"]:
--> 449 res = shell.run_cell(
450 code,
451 store_history=store_history,
452 silent=silent,
453 cell_id=cell_id,
454 )
455 else:
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/site-packages/ipykernel/zmqshell.py:549, in run_cell()
548 self._last_traceback = None
--> 549 return super().run_cell(*args, **kwargs)
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/site-packages/IPython/core/interactiveshell.py:3077, in run_cell()
3076 try:
-> 3077 result = self._run_cell(
3078 raw_cell, store_history, silent, shell_futures, cell_id
3079 )
3080 finally:
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/site-packages/IPython/core/interactiveshell.py:3132, in _run_cell()
3131 try:
-> 3132 result = runner(coro)
3133 except BaseException as e:
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/site-packages/IPython/core/async_helpers.py:128, in _pseudo_sync_runner()
127 try:
--> 128 coro.send(None)
129 except StopIteration as exc:
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/site-packages/IPython/core/interactiveshell.py:3336, in run_cell_async()
3333 interactivity = "none" if silent else self.ast_node_interactivity
-> 3336 has_raised = await self.run_ast_nodes(code_ast.body, cell_name,
3337 interactivity=interactivity, compiler=compiler, result=result)
3339 self.last_execution_succeeded = not has_raised
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/site-packages/IPython/core/interactiveshell.py:3519, in run_ast_nodes()
3518 asy = compare(code)
-> 3519 if await self.run_code(code, result, async_=asy):
3520 return True
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/site-packages/IPython/core/interactiveshell.py:3579, in run_code()
3578 else:
-> 3579 exec(code_obj, self.user_global_ns, self.user_ns)
3580 finally:
3581 # Reset our crash handler in place
Cell In[12], line 1
----> 1 get_ipython().run_line_magic('runfile', "'/disk6001/home/mq70fene/ppp-master/finite strain/bleyer/dolfinx_materials/demos/jax/finite_strain_elastoplasticity/finite_strain_elastoplasticity.py' --wdir")
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/site-packages/IPython/core/interactiveshell.py:2482, in run_line_magic()
2481 with self.builtin_trap:
-> 2482 result = fn(*args, **kwargs)
2484 # The code below prevents the output from being displayed
2485 # when using magics with decorator @output_can_be_silenced
2486 # when the last Python token in the expression is a ';'.
File /disk6001/home/mq70fene/ppp-master/finite strain/bleyer/dolfinx_materials/demos/jax/finite_strain_elastoplasticity/finite_strain_elastoplasticity.py:165
164 qmap.update_initial_state("be_bar", fem.Constant(domain, (1.0, 1.0, 1.0, 0, 0, 0)))
--> 165 qmap.update()
167 # ## Resolution
168 #
169 # As in the MFront demo, we define the custom nonlinear problem, the corresponding Newton solver, the PETSc Krylov solver and its Geometric Algebraic MultiGrid preconditioner.
170
171 # +
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/site-packages/dolfinx_materials/quadrature_map.py:320, in update()
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/site-packages/dolfinx_materials/material/generic.py:179, in integrate()
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/site-packages/dolfinx_materials/material/jax/jax_material.py:13, in wrapper()
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/site-packages/dolfinx_materials/jax_materials/finite_strain_elastoplasticity.py:90, in constitutive_update()
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/site-packages/dolfinx_materials/material/jax/jax_newton_solver.py:145, in solve()
144 tangent_solve = lambda g, y: _solve_linear_system(x, jax.jacfwd(g)(y), y)
--> 145 x_sol, data = jax.lax.custom_root(self.r, x, solve, tangent_solve, has_aux=True)
147 # x_sol, data = newton_solve(x, self.r, self.dr_dx, self.params)
JaxStackTraceBeforeTransformation: TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got:
primal int64[] with tangent int64[], expecting tangent ShapedArray(float0[])
The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.
--------------------
The above exception was the direct cause of the following exception:
TypeError Traceback (most recent call last)
File /disk6001/home/mq70fene/ppp-master/finite strain/bleyer/dolfinx_materials/demos/jax/finite_strain_elastoplasticity/finite_strain_elastoplasticity.py:165
163 qmap.initialize_state()
164 qmap.update_initial_state("be_bar", fem.Constant(domain, (1.0, 1.0, 1.0, 0, 0, 0)))
--> 165 qmap.update()
167 # ## Resolution
168 #
169 # As in the MFront demo, we define the custom nonlinear problem, the corresponding Newton solver, the PETSc Krylov solver and its Geometric Algebraic MultiGrid preconditioner.
170
171 # +
172 problem = NonlinearMaterialProblem(qmap, Res, Jac, u, bcs)
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/site-packages/dolfinx_materials/quadrature_map.py:320, in QuadratureMap.update(self)
317 Ct_vals = np.zeros_like(get_vals(self.jacobian_flatten)[self.dofs])
319 # material integration
--> 320 flux_vals, isv_vals, Ct_vals = self.material.integrate(grad_vals)
321 assert not (np.any(np.isnan(flux_vals)))
322 assert not (np.any(np.isnan(isv_vals)))
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/site-packages/dolfinx_materials/material/generic.py:179, in Material.integrate(self, gradients, dt)
176 def integrate(self, gradients, dt=0):
177 vectorized_state = self.get_initial_state_dict()
--> 179 Ct_array, new_state = self.batched_constitutive_update(
180 gradients, vectorized_state, dt
181 )
183 self.data_manager.s1.set_item(new_state)
185 return (
186 self.data_manager.s1.fluxes,
187 self.data_manager.s1.internal_state_variables,
188 Ct_array,
189 )
[... skipping hidden 20 frame]
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/site-packages/dolfinx_materials/material/jax/jax_material.py:13, in tangent_AD.<locals>.wrapper(self, *args)
10 @wraps(constitutive_update_method)
11 def wrapper(self, *args):
12 constitutive_update_fun = partial(constitutive_update_method, self)
---> 13 return jax.jacfwd(constitutive_update_fun, argnums=0, has_aux=True)(*args)
[... skipping hidden 15 frame]
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/site-packages/dolfinx_materials/jax_materials/finite_strain_elastoplasticity.py:90, in FeFpJ2Plasticity.constitutive_update(self, Fv, state, dt)
88 x0 = jnp.zeros((7,))
89 x0 = x0.at[:-1].set(to_vect(be_bar_trial, True))
---> 90 x, data = self.newton_solver.solve(x0)
92 be_bar = to_mat(x[:-1])
93 dp = x[-1]
[... skipping hidden 31 frame]
File /disk6001/home/mq70fene/anaconda3/envs/fenicsx-env/lib/python3.13/site-packages/jax/_src/custom_derivatives.py:360, in _flatten_jvp(f, store, primal_name, jvp_name, in_tree, maybe_out_type, *args)
353 msg = ("Custom JVP rule must produce primal and tangent outputs with "
354 "corresponding shapes and dtypes, but got:\n{}")
355 disagreements = (
356 f" primal {av_p.str_short()} with tangent {av_t.str_short()}, expecting tangent {av_et}"
357 for av_p, av_et, av_t in zip(primal_avals_out, expected_tangent_avals_out, tangent_avals_out)
358 if av_et != av_t)
--> 360 raise TypeError(msg.format('\n'.join(disagreements)))
361 store.store((out_tree, primal_avals))
362 return primals_out + tangents_out
TypeError: Custom JVP rule must produce primal and tangent outputs with corresponding shapes and dtypes, but got:
primal int64[] with tangent int64[], expecting tangent ShapedArray(float0[])
Thanks for the help!