fix(deps): update dependency jax to >=0.9,<0.10 #84

renovate · 2026-01-21T00:29:05Z

This PR contains the following updates:

Package	Change	Age	Confidence
jax	`>=0.8,<0.9` → `>=0.9,<0.10`

Release Notes

jax-ml/jax (jax)

`v0.9.0`

Compare Source

New features:
- Added {func}jax.thread_guard, a context manager that detects when devices
  are used by multiple threads in multi-controller JAX.
Bug fixes:
- Fixed a workspace size calculation error for pivoted QR (magma_zgeqp3_gpu)
  in MAGMA 2.9.0 when using use_magma=True and pivoting=True.
  ({jax-issue}#34145).
Deprecations:
- The flag jax_collectives_common_channel_id was removed.
- The jax_pmap_no_rank_reduction config state has been removed. The
  no-rank-reduction behavior is now the only supported behavior: a
  jax.pmapped function f sees inputs of the same rank as the input to
  jax.pmap(f). For example, if jax.pmap(f) receives shape (8, 128) on
  8 devices, then f receives shape (1, 128).
- Setting the jax_pmap_shmap_merge config state is deprecated in JAX v0.9.0
  and will be removed in JAX v0.10.0.
- {func}jax.numpy.fix is deprecated, anticipating the deprecation of
  {func}numpy.fix in NumPy v2.5.0. {func}jax.numpy.trunc is a drop-in
  replacement.
Changes:
- {func}jax.export now supports explicit sharding. This required a new
  export serialization format version that includes the NamedSharding,
  including the abstract mesh, and the partition spec. As part of this
  change we have added a restriction in the use of exported modules: when
  calling them the abstract mesh must match the one used at export time,
  including the axis names. Previously, only the number of the devices
  mattered.

Configuration

📅 Schedule: Branch creation - At any time (no schedule defined), Automerge - At any time (no schedule defined).

🚦 Automerge: Enabled.

♻ Rebasing: Whenever PR is behind base branch, or you tick the rebase/retry checkbox.

🔕 Ignore: Close this PR and you won't be reminded about these updates again.

If you want to rebase/retry this PR, check this box

This PR was generated by Mend Renovate. View the repository job log.

renovate · 2026-01-21T00:29:07Z

⚠️ Artifact update problem

Renovate failed to update an artifact related to this branch. You probably do not want to merge this PR as-is.

♻ Renovate will retry this branch, including artifacts, only when one of the following happens:

any of the package files in this branch needs updating, or
the branch becomes conflicted, or
you click the rebase/retry checkbox if found above, or
you rename this PR's title to start with "rebase!" to trigger it manually

The artifact failure details are included below:

File name: uv.lock

Command failed: uv lock --upgrade-package jax --upgrade-package jax --upgrade-package jax
Downloading cpython-3.13.11-linux-x86_64-gnu (download) (33.4MiB)
 Downloaded cpython-3.13.11-linux-x86_64-gnu (download)
Using CPython 3.13.11
  × No solution found when resolving dependencies for split (markers:
  │ python_full_version >= '3.13'):
  ╰─▶ Because liblaf-melon>=1.0.0 depends on jax>=0.8,<0.9 and only the
      following versions of liblaf-melon are available:
          liblaf-melon<=1.0.0
          liblaf-melon==1.1.0
      we can conclude that liblaf-melon>=1.0.0 depends on jax>=0.8,<0.9.
      And because your project depends on jax>=0.9, we can conclude that your
      project and liblaf-melon>=1.0.0 are incompatible.
      And because your project depends on liblaf-melon>=1 and your project
      requires liblaf-apple[cuda12], we can conclude that your project's
      requirements are unsatisfiable.

for more information, see https://pre-commit.ci

liblaf · 2026-01-21T00:30:12Z

✅⚠️MegaLinter analysis: Success with warnings

Descriptor	Linter	Files	Fixed	Errors	Warnings	Elapsed time
✅ BASH	shellcheck	2		0	0	0.03s
✅ BASH	shfmt	2	0	0	0	0.67s
⚠️ COPYPASTE	jscpd	yes		275	no	10.29s
✅ JSON	prettier	7	0	0	0	1.58s
✅ JSON	v8r	7		0	0	5.8s
✅ PYTHON	ruff	yes	yes	no	no	0.54s
✅ REPOSITORY	git_diff	yes		no	no	0.91s
⚠️ SPELL	cspell	214		9	0	8.59s
✅ YAML	prettier	4	0	0	0	1.52s
✅ YAML	v8r	4		0	0	6.33s
✅ YAML	yamllint	4		0	0	0.82s

Detailed Issues

⚠️ SPELL / cspell - 9 errors

exp/2025/10/22/inverse-flame/src/11-gen-manual.py:46:46     - Unknown word (nasi)       -- labii_superioris_alaeque_nasi001_00", 20.0),
	 Suggestions: [ansi, nasa, nash, nazi, nisi]
exp/2025/10/22/inverse-flame/src/11-gen-manual.py:47:46     - Unknown word (nasi)       -- labii_superioris_alaeque_nasi001_01", 20.0),
	 Suggestions: [ansi, nasa, nash, nazi, nisi]
exp/2025/10/22/inverse-flame/src/11-gen-manual.py:50:13     - Unknown word (Risorius)   -- # ("Risorius001_00", 10.0),
	 Suggestions: [Rigorous, roscius, Roscius, Risus, Rigors]
exp/2025/10/22/inverse-flame/src/11-gen-manual.py:51:13     - Unknown word (Risorius)   -- # ("Risorius001_01", 10.0),
	 Suggestions: [Rigorous, roscius, Roscius, Risus, Rigors]
exp/2025/10/22/inverse-flame/src/11-gen-manual.py:58:44     - Unknown word (nasi)       -- labii_superioris_alaeque_nasi001_00", 20.0),
	 Suggestions: [ansi, nasa, nash, nazi, nisi]
exp/2025/10/22/inverse-flame/src/11-gen-manual.py:59:44     - Unknown word (nasi)       -- labii_superioris_alaeque_nasi001_01", 20.0),
	 Suggestions: [ansi, nasa, nash, nazi, nisi]
exp/2025/10/22/inverse-flame/src/11-gen-manual.py:62:11     - Unknown word (Risorius)   -- ("Risorius001_00", 10.0),
	 Suggestions: [Rigorous, roscius, Roscius, Risus, Rigors]
exp/2025/10/22/inverse-flame/src/11-gen-manual.py:63:11     - Unknown word (Risorius)   -- ("Risorius001_01", 10.0),
	 Suggestions: [Rigorous, roscius, Roscius, Risus, Rigors]
tmp.6NE8UxMvCk/renovate-config.json:28:4      - Unknown word (pyenv)      -- "pyenv": {
	 Suggestions: [peen, pena, pend, peng, penh]
CSpell: Files checked: 214, Issues found: 9 in 2 files.


You can skip this misspellings by defining the following .cspell.json file at the root of your repository
Of course, please correct real typos before :)

{
    "version": "0.2",
    "language": "en",
    "ignorePaths": [
        "**/node_modules/**",
        "**/vscode-extension/**",
        "**/.git/**",
        "**/.pnpm-lock.json",
        ".vscode",
        "package-lock.json",
        "megalinter-reports"
    ],
    "words": [
        "Risorius",
        "nasi",
        "pyenv"
    ]
}


You can also copy-paste megalinter-reports/LINTER_DEFAULT at the root of your repository

⚠️ COPYPASTE / jscpd - 275 errors

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_phace_fixed_hess.py [127:6 - 133:18] (6 lines, 88 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_phace_fixed_hess.py [104:6 - 110:18]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_phace.py [1:1 - 23:6] (22 lines, 171 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_phace_fixed_hess.py [1:1 - 23:13]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_phace.py [23:6 - 45:6] (22 lines, 168 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_phace_fixed_hess.py [23:13 - 45:13]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_phace.py [45:6 - 61:6] (16 lines, 142 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_phace_fixed_hess.py [45:13 - 61:13]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_phace.py [64:6 - 86:6] (22 lines, 217 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_phace_fixed_hess.py [64:13 - 86:13]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_phace.py [89:6 - 110:17] (21 lines, 220 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_phace_fixed_hess.py [89:13 - 110:18]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_phace.py [119:3 - 139:17] (20 lines, 240 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_phace_fixed_hess.py [113:3 - 110:18]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_phace.py [148:3 - 168:17] (20 lines, 246 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_phace_fixed_hess.py [136:3 - 156:18]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_phace.py [178:9 - 197:15] (19 lines, 249 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_phace_fixed_hess.py [160:9 - 179:16]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_phace.py [201:6 - 207:3] (6 lines, 85 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_phace_fixed_hess.py [180:13 - 186:7]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_phace.py [223:2 - 250:7] (27 lines, 294 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_phace_fixed_hess.py [190:13 - 215:7]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py [222:12 - 231:5] (9 lines, 101 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py [197:11 - 206:7]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py [231:5 - 242:45] (11 lines, 162 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py [206:7 - 217:7]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py [255:17 - 275:10] (20 lines, 264 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py [197:11 - 217:7]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py [297:2 - 312:3] (15 lines, 219 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py [201:2 - 216:12]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py [331:17 - 357:10] (26 lines, 361 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py [292:17 - 318:10]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py [372:23 - 381:16] (9 lines, 92 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py [197:11 - 206:7]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py [381:6 - 392:21] (11 lines, 162 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py [206:7 - 217:7]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py [404:22 - 413:6] (9 lines, 92 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py [197:11 - 206:7]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py [414:5 - 425:6] (11 lines, 171 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py [231:7 - 217:7]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py [439:26 - 448:5] (9 lines, 92 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py [197:11 - 206:7]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py [449:10 - 464:9] (15 lines, 200 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_base.py [231:7 - 246:5]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle_v2.py [87:6 - 92:8] (5 lines, 85 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle_v2.py [72:6 - 77:8]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle_v2.py [131:9 - 139:7] (8 lines, 108 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_phace_fixed_hess.py [203:3 - 211:7]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle.py [1:1 - 21:11] (20 lines, 147 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle_v2.py [1:1 - 21:13]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle.py [21:11 - 39:11] (18 lines, 122 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle_v2.py [21:13 - 39:13]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle.py [39:11 - 51:2] (12 lines, 134 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle_v2.py [39:13 - 51:2]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle.py [51:2 - 64:2] (13 lines, 121 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle_v2.py [51:2 - 64:2]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle.py [65:9 - 76:7] (11 lines, 96 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle_v2.py [65:9 - 76:2]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle.py [79:7 - 92:7] (13 lines, 147 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle_v2.py [78:5 - 91:2]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle.py [88:6 - 94:8] (6 lines, 98 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle.py [72:6 - 78:8]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle.py [95:7 - 108:7] (13 lines, 147 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle_v2.py [93:5 - 106:2]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle.py [111:7 - 124:7] (13 lines, 149 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle_v2.py [108:5 - 121:2]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle.py [127:7 - 143:7] (16 lines, 245 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle_v2.py [123:5 - 139:7]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_active.py [160:11 - 181:18] (21 lines, 215 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_phace_fixed_hess.py [190:6 - 211:10]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap.py [69:2 - 81:2] (12 lines, 138 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle_v2.py [77:2 - 89:2]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap.py [82:2 - 94:2] (12 lines, 141 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle_v2.py [92:2 - 104:2]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap.py [95:2 - 107:2] (12 lines, 143 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle_v2.py [107:2 - 119:2]

Clone found (python):
 - src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap.py [108:2 - 123:5] (15 lines, 213 tokens)
   src/liblaf/apple/warp/energies/elastic/hyperelastic/_arap_muscle_v2.py [122:2 - 137:13]

Clone found (python):
 - exp/2025/07/30/dynamics/collision/collision-bunny-bunny/main.py [1:1 - 14:5] (13 lines, 115 tokens)
   exp/2025/07/30/dynamics/collision/collision-bunny-sphere/main.py [1:1 - 14:4]

Clone found (python):
 - exp/2025/07/30/dynamics/collision/collision-bunny-bunny/main.py [22:5 - 84:78] (62 lines, 768 tokens)
   exp/2025/07/30/dynamics/collision/collision-bunny-sphere/main.py [22:5 - 84:8]

Clone found (python):
 - exp/2025/07/30/dynamics/collision/collision-bunny-bunny/main.py [87:5 - 114:2] (27 lines, 319 tokens)
   exp/2025/07/30/dynamics/collision/collision-bunny-sphere/main.py [84:4 - 111:2]

Clone found (python):
 - exp/2025/07/30/dynamics/collision/collision-bunny/main.py [21:5 - 35:7] (14 lines, 116 tokens)
   exp/2025/07/30/dynamics/collision/collision-bunny-bunny/main.py [17:5 - 31:4]

Clone found (python):
 - exp/2025/07/30/dynamics/collision/collision-bunny/main.py [58:1 - 67:4] (9 lines, 151 tokens)
   exp/2025/07/30/dynamics/collision/collision-bunny-sphere/main.py [31:1 - 40:3]

Clone found (python):
 - exp/2025/07/30/dynamics/collision/collision-bunny/main.py [67:3 - 77:9] (10 lines, 124 tokens)
   exp/2025/07/30/dynamics/collision/collision-bunny-sphere/main.py [40:3 - 49:6]

Clone found (python):
 - exp/2025/07/30/dynamics/collision/collision-bunny/main.py [81:9 - 88:5] (7 lines, 101 tokens)
   exp/2025/07/30/dynamics/collision/collision-bunny-sphere/main.py [52:9 - 59:7]

Clone found (python):
 - exp/2025/07/30/dynamics/collision/collision-bunny/main.py [90:4 - 100:5] (10 lines, 128 tokens)
   exp/2025/07/30/dynamics/collision/collision-bunny-sphere/main.py [58:2 - 68:4]

Clone found (python):
 - exp/2025/07/30/dynamics/collision/collision-bunny/main.py [100:2 - 116:4] (16 lines, 181 tokens)
   exp/2025/07/30/dynamics/collision/collision-bunny-sphere/main.py [68:2 - 84:10]

Clone found (python):
 - exp/2025/07/30/dynamics/collision/collision-bunny/main.py [116:6 - 132:5] (16 lines, 200 tokens)
   exp/2025/07/30/dynamics/collision/collision-bunny-sphere/main.py [87:5 - 103:13]

Clone found (python):
 - exp/2025/07/30/dynamics/collision/collision-box/main.py [1:1 - 50:5] (49 lines, 490 tokens)
   exp/2025/07/30/dynamics/collision/collision-bunny/main.py [1:1 - 50:3]

Clone found (python):
 - exp/2025/07/30/dynamics/collision/collision-box/main.py [52:9 - 79:28] (27 lines, 359 tokens)
   exp/2025/07/30/dynamics/collision/collision-bunny/main.py [54:2 - 81:3]

Clone found (python):
 - exp/2025/07/30/dynamics/collision/collision-box/main.py [78:9 - 96:75] (18 lines, 224 tokens)
   exp/2025/07/30/dynamics/collision/collision-bunny-sphere/main.py [49:10 - 96:8]

Clone found (python):
 - exp/2025/07/30/dynamics/collision/collision-box/main.py [104:5 - 142:10] (38 lines, 481 tokens)
   exp/2025/07/30/dynamics/collision/collision-bunny/main.py [97:8 - 105:12]

Clone found (python):
 - tests/warp/energies/elastic/hyperelastic/func/test_hess_quad.py [50:6 - 60:2] (10 lines, 99 tokens)
   tests/warp/energies/elastic/hyperelastic/func/test_hess_quad.py [24:1 - 34:7]

Clone found (python):
 - tests/warp/energies/elastic/hyperelastic/func/test_hess_quad.py [64:5 - 73:8] (9 lines, 160 tokens)
   tests/warp/energies/elastic/hyperelastic/func/test_hess_quad.py [41:5 - 50:8]

Clone found (python):
 - tests/warp/energies/elastic/hyperelastic/func/test_hess_quad.py [77:3 - 83:8] (6 lines, 75 tokens)
   tests/warp/energies/elastic/hyperelastic/func/test_hess_quad.py [28:2 - 34:7]

Clone found (python):
 - tests/warp/energies/elastic/hyperelastic/func/test_hess_quad.py [86:5 - 95:8] (9 lines, 160 tokens)
   tests/warp/energies/elastic/hyperelastic/func/test_hess_quad.py [41:5 - 50:8]

Clone found (python):
 - tests/warp/energies/elastic/hyperelastic/func/test_hess_quad.py [95:8 - 105:8] (10 lines, 116 tokens)
   tests/warp/energies/elastic/hyperelastic/func/test_hess_quad.py [50:8 - 34:7]

Clone found (python):
 - tests/warp/energies/elastic/hyperelastic/func/test_hess_quad.py [108:5 - 117:12] (9 lines, 158 tokens)
   tests/warp/energies/elastic/hyperelastic/func/test_hess_quad.py [41:5 - 50:4]

Clone found (python):
 - tests/warp/energies/elastic/hyperelastic/func/test_hess_quad.py [141:3 - 150:13] (9 lines, 90 tokens)
   tests/warp/energies/elastic/hyperelastic/func/test_hess_quad.py [130:3 - 139:13]

Clone found (python):
 - tests/warp/energies/elastic/hyperelastic/func/test_hess_prod.py [99:8 - 109:22] (10 lines, 97 tokens)
   tests/warp/energies/elastic/hyperelastic/func/test_hess_prod.py [76:8 - 86:13]

Clone found (python):
 - tests/warp/energies/elastic/hyperelastic/func/test_hess_prod.py [109:22 - 120:6] (11 lines, 182 tokens)
   tests/warp/energies/elastic/hyperelastic/func/test_hess_prod.py [86:13 - 97:5]

Clone found (python):
 - tests/warp/energies/elastic/hyperelastic/func/test_hess_prod.py [120:6 - 125:2] (5 lines, 93 tokens)
   tests/warp/energies/elastic/hyperelastic/func/test_hess_prod.py [97:5 - 104:2]

Clone found (python):
 - tests/warp/energies/elastic/hyperelastic/func/test_hess_diag.py [68:5 - 73:13] (5 lines, 65 tokens)
   tests/warp/energies/elastic/hyperelastic/func/test_hess_prod.py [81:5 - 86:13]

Clone found (python):
 - tests/warp/energies/elastic/hyperelastic/func/test_hess_diag.py [74:5 - 85:11] (11 lines, 208 tokens)
   tests/warp/energies/elastic/hyperelastic/func/test_hess_prod.py [88:5 - 99:11]

Clone found (python):
 - tests/warp/energies/elastic/hyperelastic/func/test_hess_diag.py [85:8 - 94:22] (9 lines, 83 tokens)
   tests/warp/energies/elastic/hyperelastic/func/test_hess_diag.py [64:8 - 86:13]

Clone found (python):
 - tests/warp/energies/elastic/hyperelastic/func/test_hess_diag.py [94:22 - 106:11] (12 lines, 211 tokens)
   tests/warp/energies/elastic/hyperelastic/func/test_hess_diag.py [73:13 - 122:11]

Clone found (python):
 - tests/warp/energies/elastic/hyperelastic/func/test_hess_diag.py [104:6 - 109:2] (5 lines, 80 tokens)
   tests/warp/energies/elastic/hyperelastic/func/test_hess_diag.py [83:5 - 89:2]

Clone found (python):
 - exp/2025/12/31/inverse-toy/src/20-forward-newton.py [7:1 - 31:15] (24 lines, 232 tokens)
   exp/2025/12/31/inverse-toy/src/20-forward.py [7:1 - 32:5]

Clone found (python):
 - exp/2025/12/31/inverse-toy/src/20-forward-newton.py [32:5 - 53:32] (21 lines, 280 tokens)
   exp/2025/12/31/inverse-toy/src/20-forward.py [39:5 - 60:25]

Clone found (python):
 - exp/2025/12/31/inverse-toy/src/20-forward-muscle.py [33:1 - 43:5] (10 lines, 109 tokens)
   exp/2025/12/31/inverse-toy/src/20-forward.py [24:1 - 34:5]

Clone found (python):
 - exp/2025/12/31/inverse-toy/src/20-forward-muscle.py [62:5 - 67:2] (5 lines, 79 tokens)
   exp/2025/12/31/inverse-toy/src/20-forward.py [51:5 - 56:4]

Clone found (python):
 - exp/2025/10/22/inverse-flame/src/21-inverse-inspect.py [28:5 - 33:2] (5 lines, 110 tokens)
   exp/2025/10/22/inverse-flame/src/31-animate-inspect.py [27:13 - 32:2]

Clone found (python):
 - exp/2025/10/22/inverse-flame/src/21-inverse-inspect.py [46:5 - 57:2] (11 lines, 184 tokens)
   exp/2025/10/22/inverse-flame/src/31-animate-inspect.py [54:13 - 65:2]

Clone found (python):
 - exp/2025/10/22/inverse-flame/src/20-inverse-adam.py [19:1 - 31:7] (12 lines, 136 tokens)
   exp/2025/10/22/inverse-flame/src/20-inverse-lbfgs.py [16:1 - 27:7]

Clone found (python):
 - exp/2025/10/22/inverse-flame/src/20-inverse-adam.py [60:9 - 68:18] (8 lines, 117 tokens)
   exp/2025/10/22/inverse-flame/src/20-inverse-lbfgs.py [51:9 - 59:28]

Clone found (python):
 - exp/2025/10/22/inverse-flame/src/20-inverse-adam.py [108:9 - 124:2] (16 lines, 178 tokens)
   exp/2025/10/22/inverse-flame/src/20-inverse-lbfgs.py [83:9 - 99:6]

Clone found (python):
 - exp/2025/10/22/inverse-flame/src/20-inverse-adam.py [141:5 - 151:15] (10 lines, 150 tokens)
   exp/2025/10/22/inverse-flame/src/20-inverse-lbfgs.py [107:5 - 117:14]

Clone found (python):
 - exp/2025/10/22/inverse-flame/src/20-inverse-adam.py [163:15 - 172:10] (9 lines, 97 tokens)
   exp/2025/10/22/inverse-flame/src/20-inverse-lbfgs.py [133:14 - 143:4]

Clone found (python):
 - exp/2025/10/22/inverse-flame/src/20-inverse-adam.py [190:1 - 218:5] (28 lines, 306 tokens)
   exp/2025/10/22/inverse-flame/src/20-inverse-lbfgs.py [143:1 - 171:8]

Clone found (python):
 - exp/2025/10/22/inverse-flame/src/20-inverse-adam.py [230:5 - 245:15] (15 lines, 186 tokens)
   exp/2025/10/22/inverse-flame/src/20-inverse-lbfgs.py [175:5 - 190:28]

Clone found (python):
 - exp/2025/10/22/inverse-flame/src/20-inverse-adam.py [263:2 - 277:5] (14 lines, 118 tokens)
   exp/2025/10/22/inverse-flame/src/20-inverse-lbfgs.py [203:2 - 217:8]

Clone found (python):
 - exp/2025/10/22/inverse-flame/src/20-inverse-adam.py [282:7 - 295:19] (13 lines, 180 tokens)
   exp/2025/10/22/inverse-flame/src/20-inverse-lbfgs.py [224:6 - 236:9]

Clone found (python):
 - exp/2025/10/22/inverse-flame/src/20-inverse-adam-stage-2.py [6:1 - 34:6] (28 lines, 294 tokens)
   exp/2025/10/22/inverse-flame/src/20-inverse-adam.py [8:1 - 36:11]

Clone found (python):
 - exp/2025/10/22/inverse-flame/src/20-inverse-adam-stage-2.py [41:1 - 176:3] (135 lines, 1531 tokens)
   exp/2025/10/22/inverse-flame/src/20-inverse-adam.py [42:1 - 177:9]

Clone found (python):
 - exp/2025/10/22/inverse-flame/src/20-inverse-adam-stage-2.py [179:9 - 275:8] (96 lines, 996 tokens)
   exp/2025/10/22/inverse-flame/src/20-inverse-adam.py [176:9 - 272:8]

Clone found (python):
 - exp/2025/10/22/inverse-flame/src/20-inverse-adam-stage-2.py [283:45 - 325:75] (42 lines, 430 tokens)
   exp/2025/10/22/inverse-flame/src/20-inverse-adam.py [278:58 - 320:4]

Clone found (python):
 - exp/2025/10/22/inverse-flame/src/20-inverse-adam-partial.py [1:1 - 30:8] (29 lines, 309 tokens)
   exp/2025/10/22/inverse-flame/src/20-inverse-adam-stage-2.py [1:1 - 32:8]

Clone found (python):
 - exp/2025/10/22/inverse-flame/src/20-inverse-adam-partial.py [37:31 - 277:66] (240 lines, 2597 tokens)
   exp/2025/10/22/inverse-flame/src/20-inverse-adam.py [39:23 - 278:58]

Clone found (python):
 - exp/2025/10/22/inverse-flame/src/20-inverse-adam-partial.py [278:9 - 320:8] (42 lines, 430 tokens)
   exp/2025/10/22/inverse-flame/src/20-inverse-adam.py [278:58 - 320:4]

Clone found (python):
 - exp/2025/10/22/inverse-flame/src/20-inverse-adam-old-hess.py [1:1 - 35:11] (34 lines, 341 tokens)
   exp/2025/10/22/inverse-flame/src/20-inverse-adam-stage-2.py [1:1 - 35:8]

Clone found (python):
 - exp/2025/10/22/inverse-flame/src/20-inverse-adam-old-hess.py [35:16 - 274:63] (239 lines, 2595 tokens)
   exp/2025/10/22/inverse-flame/src/20-inverse-adam.py [39:23 - 278:58]

Clone found (python):
 - exp/2025/10/22/inverse-flame/src/20-inverse-adam-old-hess.py [274:63 - 317:4] (43 lines, 469 tokens)
   exp/2025/10/22/inverse-flame/src/20-inverse-adam.py [278:58 - 321:3]

Clone found (python):
 - exp/2025/10/22/inverse-flame/src/20-inverse-adam-old-hess.py [316:4 - 330:2] (14 lines, 88 tokens)
   exp/2025/10/22/inverse-flame/src/20-inverse-adam.py [323:5 - 337:2]

Clone found (python):
 - exp/2025/10/22/inverse-flame/src/20-inverse-adam-fine.py [1:1 - 34:6] (33 lines, 325 tokens)
   exp/2025/10/22/inverse-flame/src/20-inverse-adam-stage-2.py [1:1 - 34:11]

Clone found (python):
 - exp/2025/10/22/inverse-flame/src/20-inverse-adam-fine.py [37:5 - 54:2] (17 lines, 101 tokens)
   exp/2025/10/22/inverse-flame/src/20-inverse-adam-stage-2.py [37:5 - 57:6]

Clone found (python):
 - exp/2025/10/22/inverse-flame/src/20-inverse-adam-fine.py [56:5 - 272:2] (216 lines, 2456 tokens)
   exp/2025/10/22/inverse-flame/src/20-inverse-adam.py [57:5 - 276:2]

Clone found (python):
 - exp/2025/10/22/inverse-

(Truncated to 20000 characters out of 52957)

See detailed reports in MegaLinter artifacts

codecov · 2026-01-21T01:40:13Z

❌ 12 Tests Failed:

Tests completed	Failed	Passed	Skipped
59	12	47	1

View the full list of 12 ❄️ flaky test(s)

tests/jax/math/test_rotation.py::test_polar_rv

Flake rate in main: 100.00% (Passed 0 times, Failed 15 times)

Stack Traces | 8.93s run time

#x1B[0m#x1B[37m@hypothesis#x1B[39;49;00m.given(testing.matrices((#x1B[94m3#x1B[39;49;00m, #x1B[94m3#x1B[39;49;00m)))#x1B[90m#x1B[39;49;00m
>   #x1B[94mdef#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[92mtest_polar_rv#x1B[39;49;00m(F: Mat33) -> #x1B[94mNone#x1B[39;49;00m:#x1B[90m#x1B[39;49;00m
                   ^^^#x1B[90m#x1B[39;49;00m

f          = <function given.<locals>.run_test_as_given.<locals>.wrapped_test at 0x7f3c22d62160>

#x1B[1m#x1B[.../jax/math/test_rotation.py#x1B[0m:33: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

F = Array([[[-0.76369951, -0.09248251,  0.13710512],
        [ 0.66158777, -0.59512941,  0.51031201],
        [ 0.06014349, -0.55544681, -0.04039191]]], dtype=float64)

    #x1B[0m#x1B[37m@hypothesis#x1B[39;49;00m.given(testing.matrices((#x1B[94m3#x1B[39;49;00m, #x1B[94m3#x1B[39;49;00m)))#x1B[90m#x1B[39;49;00m
    #x1B[94mdef#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[92mtest_polar_rv#x1B[39;49;00m(F: Mat33) -> #x1B[94mNone#x1B[39;49;00m:#x1B[90m#x1B[39;49;00m
        R: Mat33#x1B[90m#x1B[39;49;00m
        S: Mat33#x1B[90m#x1B[39;49;00m
        R, S = math.polar_rv(F)#x1B[90m#x1B[39;49;00m
        np.testing.assert_allclose(#x1B[90m#x1B[39;49;00m
            R.mT @ R, jnp.broadcast_to(jnp.identity(#x1B[94m3#x1B[39;49;00m), F.shape), atol=ATOL#x1B[90m#x1B[39;49;00m
        )#x1B[90m#x1B[39;49;00m
        np.testing.assert_allclose(jnp.linalg.det(R), #x1B[94m1.0#x1B[39;49;00m)#x1B[90m#x1B[39;49;00m
>       np.testing.assert_allclose(R @ S, F, atol=ATOL)#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31mE       AssertionError: #x1B[0m
#x1B[1m#x1B[31mE       Not equal to tolerance rtol=1e-07, atol=1e-07#x1B[0m
#x1B[1m#x1B[31mE       #x1B[0m
#x1B[1m#x1B[31mE       Mismatched elements: 9 / 9 (100%)#x1B[0m
#x1B[1m#x1B[31mE       First 5 mismatches are at indices:#x1B[0m
#x1B[1m#x1B[31mE        [0, 0, 0]: -0.7507176526526799 (ACTUAL), -0.7636995066742065 (DESIRED)#x1B[0m
#x1B[1m#x1B[31mE        [0, 0, 1]: -0.20091949436875461 (ACTUAL), -0.09248251142607344 (DESIRED)#x1B[0m
#x1B[1m#x1B[31mE        [0, 0, 2]: -0.08149924494380247 (ACTUAL), 0.13710511563780203 (DESIRED)#x1B[0m
#x1B[1m#x1B[31mE        [0, 1, 0]: 0.6782985317709115 (ACTUAL), 0.6615877671366968 (DESIRED)#x1B[0m
#x1B[1m#x1B[31mE        [0, 1, 1]: -0.7347138547914873 (ACTUAL), -0.5951294089492527 (DESIRED)#x1B[0m
#x1B[1m#x1B[31mE       Max absolute difference among violations: 0.51135173#x1B[0m
#x1B[1m#x1B[31mE       Max relative difference among violations: 12.65975723#x1B[0m
#x1B[1m#x1B[31mE        ACTUAL: array([[[-0.750718, -0.200919, -0.081499],#x1B[0m
#x1B[1m#x1B[31mE               [ 0.678299, -0.734714,  0.228916],#x1B[0m
#x1B[1m#x1B[31mE               [ 0.029777, -0.301795,  0.47096 ]]])#x1B[0m
#x1B[1m#x1B[31mE        DESIRED: array([[[-0.7637  , -0.092483,  0.137105],#x1B[0m
#x1B[1m#x1B[31mE               [ 0.661588, -0.595129,  0.510312],#x1B[0m
#x1B[1m#x1B[31mE               [ 0.060143, -0.555447, -0.040392]]])#x1B[0m
#x1B[1m#x1B[31mE       Falsifying example: test_polar_rv(#x1B[0m
#x1B[1m#x1B[31mE           F=Array([[[-0.76369951, -0.09248251,  0.13710512],#x1B[0m
#x1B[1m#x1B[31mE                   [ 0.66158777, -0.59512941,  0.51031201],#x1B[0m
#x1B[1m#x1B[31mE                   [ 0.06014349, -0.55544681, -0.04039191]]], dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE       )#x1B[0m
#x1B[1m#x1B[31mE       #x1B[0m
#x1B[1m#x1B[31mE       You can reproduce this example by temporarily adding @reproduce_failure('6.150.1', b'AEEBQQE=') as a decorator on your test case#x1B[0m

F          = Array([[[-0.76369951, -0.09248251,  0.13710512],
        [ 0.66158777, -0.59512941,  0.51031201],
        [ 0.06014349, -0.55544681, -0.04039191]]], dtype=float64)
R          = Array([[[-0.83508667, -0.51821654, -0.18461274],
        [ 0.53520329, -0.84294523, -0.05477937],
        [-0.12723085, -0.14455087,  0.98128352]]], dtype=float64)
S          = Array([[[ 0.98615338, -0.18703847,  0.13065474],
        [-0.18703847,  0.76706804, -0.21880678],
        [ 0.13065474, -0.21880678,  0.46465105]]], dtype=float64)

#x1B[1m#x1B[.../jax/math/test_rotation.py#x1B[0m:41: AssertionError

tests/jax/math/test_rotation.py::test_svd_rv

Flake rate in main: 100.00% (Passed 0 times, Failed 15 times)

Stack Traces | 12.2s run time

#x1B[0m#x1B[37m@hypothesis#x1B[39;49;00m.given(testing.matrices((#x1B[94m3#x1B[39;49;00m, #x1B[94m3#x1B[39;49;00m)))#x1B[90m#x1B[39;49;00m
>   #x1B[94mdef#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[92mtest_svd_rv#x1B[39;49;00m(F: Mat33) -> #x1B[94mNone#x1B[39;49;00m:#x1B[90m#x1B[39;49;00m
                   ^^^#x1B[90m#x1B[39;49;00m

f          = <function given.<locals>.run_test_as_given.<locals>.wrapped_test at 0x7f3c22d62840>

#x1B[1m#x1B[.../jax/math/test_rotation.py#x1B[0m:15: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

F = Array([[[-0.76369951, -0.09248251,  0.13710512],
        [ 0.66158777, -0.59512941,  0.51031201],
        [ 0.06014349, -0.55544681, -0.04039191]]], dtype=float64)

    #x1B[0m#x1B[37m@hypothesis#x1B[39;49;00m.given(testing.matrices((#x1B[94m3#x1B[39;49;00m, #x1B[94m3#x1B[39;49;00m)))#x1B[90m#x1B[39;49;00m
    #x1B[94mdef#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[92mtest_svd_rv#x1B[39;49;00m(F: Mat33) -> #x1B[94mNone#x1B[39;49;00m:#x1B[90m#x1B[39;49;00m
        U: Mat33#x1B[90m#x1B[39;49;00m
        S: Vec3#x1B[90m#x1B[39;49;00m
        Vh: Mat33#x1B[90m#x1B[39;49;00m
        U, S, Vh = math.svd_rv(F)#x1B[90m#x1B[39;49;00m
        np.testing.assert_allclose(#x1B[90m#x1B[39;49;00m
            U.mT @ U, jnp.broadcast_to(jnp.identity(#x1B[94m3#x1B[39;49;00m), F.shape), atol=ATOL#x1B[90m#x1B[39;49;00m
        )#x1B[90m#x1B[39;49;00m
        np.testing.assert_allclose(jnp.linalg.det(U), #x1B[94m1.0#x1B[39;49;00m)#x1B[90m#x1B[39;49;00m
        np.testing.assert_allclose(#x1B[90m#x1B[39;49;00m
            Vh.mT @ Vh, jnp.broadcast_to(jnp.identity(#x1B[94m3#x1B[39;49;00m), F.shape), atol=ATOL#x1B[90m#x1B[39;49;00m
        )#x1B[90m#x1B[39;49;00m
        np.testing.assert_allclose(jnp.linalg.det(Vh), #x1B[94m1.0#x1B[39;49;00m)#x1B[90m#x1B[39;49;00m
        S: Mat33 = S[..., jnp.newaxis] * jnp.identity(#x1B[94m3#x1B[39;49;00m)#x1B[90m#x1B[39;49;00m
>       np.testing.assert_allclose(U @ S @ Vh, F, atol=ATOL)#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31mE       AssertionError: #x1B[0m
#x1B[1m#x1B[31mE       Not equal to tolerance rtol=1e-07, atol=1e-07#x1B[0m
#x1B[1m#x1B[31mE       #x1B[0m
#x1B[1m#x1B[31mE       Mismatched elements: 9 / 9 (100%)#x1B[0m
#x1B[1m#x1B[31mE       First 5 mismatches are at indices:#x1B[0m
#x1B[1m#x1B[31mE        [0, 0, 0]: -0.75071765265268 (ACTUAL), -0.7636995066742065 (DESIRED)#x1B[0m
#x1B[1m#x1B[31mE        [0, 0, 1]: -0.2009194943687544 (ACTUAL), -0.09248251142607344 (DESIRED)#x1B[0m
#x1B[1m#x1B[31mE        [0, 0, 2]: -0.08149924494380253 (ACTUAL), 0.13710511563780203 (DESIRED)#x1B[0m
#x1B[1m#x1B[31mE        [0, 1, 0]: 0.6782985317709118 (ACTUAL), 0.6615877671366968 (DESIRED)#x1B[0m
#x1B[1m#x1B[31mE        [0, 1, 1]: -0.7347138547914875 (ACTUAL), -0.5951294089492527 (DESIRED)#x1B[0m
#x1B[1m#x1B[31mE       Max absolute difference among violations: 0.51135173#x1B[0m
#x1B[1m#x1B[31mE       Max relative difference among violations: 12.65975723#x1B[0m
#x1B[1m#x1B[31mE        ACTUAL: array([[[-0.750718, -0.200919, -0.081499],#x1B[0m
#x1B[1m#x1B[31mE               [ 0.678299, -0.734714,  0.228916],#x1B[0m
#x1B[1m#x1B[31mE               [ 0.029777, -0.301795,  0.47096 ]]])#x1B[0m
#x1B[1m#x1B[31mE        DESIRED: array([[[-0.7637  , -0.092483,  0.137105],#x1B[0m
#x1B[1m#x1B[31mE               [ 0.661588, -0.595129,  0.510312],#x1B[0m
#x1B[1m#x1B[31mE               [ 0.060143, -0.555447, -0.040392]]])#x1B[0m
#x1B[1m#x1B[31mE       Falsifying example: test_svd_rv(#x1B[0m
#x1B[1m#x1B[31mE           F=Array([[[-0.76369951, -0.09248251,  0.13710512],#x1B[0m
#x1B[1m#x1B[31mE                   [ 0.66158777, -0.59512941,  0.51031201],#x1B[0m
#x1B[1m#x1B[31mE                   [ 0.06014349, -0.55544681, -0.04039191]]], dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE       )#x1B[0m
#x1B[1m#x1B[31mE       Explanation:#x1B[0m
#x1B[1m#x1B[31mE           These lines were always and only run by failing examples:#x1B[0m
#x1B[1m#x1B[31mE               .../apple/apple/.venv/lib/python3.12.../numpy/_core/arrayprint.py:1020#x1B[0m
#x1B[1m#x1B[31mE               .../apple/apple/.venv/lib/python3.12.../numpy/_core/arrayprint.py:1025#x1B[0m
#x1B[1m#x1B[31mE               .../apple/apple/.venv/lib/python3.12.../numpy/_core/arrayprint.py:1065#x1B[0m
#x1B[1m#x1B[31mE               .../apple/apple/.venv/lib/python3.12.../numpy/_core/fromnumeric.py:53#x1B[0m
#x1B[1m#x1B[31mE               .../apple/apple/.venv/lib/python3.12.../numpy/_core/numeric.py:672#x1B[0m
#x1B[1m#x1B[31mE               (and 7 more with settings.verbosity >= verbose)#x1B[0m
#x1B[1m#x1B[31mE       #x1B[0m
#x1B[1m#x1B[31mE       You can reproduce this example by temporarily adding @reproduce_failure('6.150.1', b'AEEBQQE=') as a decorator on your test case#x1B[0m

F          = Array([[[-0.76369951, -0.09248251,  0.13710512],
        [ 0.66158777, -0.59512941,  0.51031201],
        [ 0.06014349, -0.55544681, -0.04039191]]], dtype=float64)
S          = Array([[[1.16687313, 0.        , 0.        ],
        [0.        , 0.70264464, 0.        ],
        [0.        , 0.        , 0.34835471]]], dtype=float64)
U          = Array([[[-0.43295697, -0.83037767, -0.35074375],
        [ 0.85503034, -0.25510504, -0.45149147],
        [ 0.28543194, -0.49537293,  0.82044761]]], dtype=float64)
Vh         = Array([[[ 0.78285589, -0.53763772,  0.313181  ],
        [ 0.619931  ,  0.71696161, -0.3188285 ],
        [-0.05312453,  0.44374738,  0.8945759 ]]], dtype=float64)

#x1B[1m#x1B[.../jax/math/test_rotation.py#x1B[0m:29: AssertionError

tests/model/test_forward.py::test_forward

Flake rate in main: 48.08% (Passed 27 times, Failed 25 times)

Stack Traces | 53s run time

#x1B[0m#x1B[94mdef#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[92mtest_forward#x1B[39;49;00m() -> #x1B[94mNone#x1B[39;49;00m:#x1B[90m#x1B[39;49;00m
        mesh: pv.UnstructuredGrid = pv.examples.download_letter_a()  #x1B[90m# pyright: ignore[reportAssignmentType]#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
        mesh.cell_data[MU] = np.full((mesh.n_cells,), #x1B[94m1.0#x1B[39;49;00m)#x1B[90m#x1B[39;49;00m
        mesh.point_data[DIRICHLET_MASK] = mesh.points[:, #x1B[94m1#x1B[39;49;00m] < mesh.bounds.y_min + #x1B[94m0.1#x1B[39;49;00m * (#x1B[90m#x1B[39;49;00m
            mesh.bounds.y_max - mesh.bounds.y_min#x1B[90m#x1B[39;49;00m
        )#x1B[90m#x1B[39;49;00m
        mesh.point_data[DIRICHLET_VALUE] = np.zeros((mesh.n_points, #x1B[94m3#x1B[39;49;00m))#x1B[90m#x1B[39;49;00m
    #x1B[90m#x1B[39;49;00m
        builder = ModelBuilder()#x1B[90m#x1B[39;49;00m
        mesh = builder.assign_global_ids(mesh)#x1B[90m#x1B[39;49;00m
        builder.add_dirichlet(mesh)#x1B[90m#x1B[39;49;00m
        elastic: Arap = Arap.from_pyvista(mesh)#x1B[90m#x1B[39;49;00m
        builder.add_energy(elastic)#x1B[90m#x1B[39;49;00m
    #x1B[90m#x1B[39;49;00m
        model: Model = builder.finalize()#x1B[90m#x1B[39;49;00m
        rng: np.random.Generator = np.random.default_rng()#x1B[90m#x1B[39;49;00m
        model.u_free = jnp.asarray(#x1B[90m#x1B[39;49;00m
            rng.uniform(-mesh.length, mesh.length, model.u_free.shape)#x1B[90m#x1B[39;49;00m
        )#x1B[90m#x1B[39;49;00m
        forward = Forward(model, optimizer=PNCG(max_steps=#x1B[94m1000#x1B[39;49;00m, rtol=#x1B[94m1e-15#x1B[39;49;00m))#x1B[90m#x1B[39;49;00m
        solution: PNCG.Solution = forward.step()#x1B[90m#x1B[39;49;00m
>       #x1B[94massert#x1B[39;49;00m solution.success#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31mE       AssertionError: assert False#x1B[0m
#x1B[1m#x1B[31mE        +  where False = OptimizeSolution(\n  result=<Result.STAGNATION: 'stagnation'>,\n  state=PNCGState(\n    structure=Structure(full_flat=f64...PNCGStats(\n    n_steps=372,\n    relative_decrease=Array(1.9753192e-11, dtype=float64),\n    time=50.61154793199967\n  )\n).success#x1B[0m

builder    = ModelBuilder(
  edges_length_sum=np.float64(414.4637687552529),
  n_edges=7119,
  dirichlet=DirichletBuilder(mask=bool...      params=Arap__Params(
               	mu=array(shape=(4802,), dtype=float64),
               )
      )
    }
  )
)
elastic    = Arap(
  id='Arap000',
  cells=array(shape=(4802,), dtype=vec4i),
  dhdX=array(shape=(4802, 1), dtype=mat43(d)),
  dV=a...(shape=(4802, 1), dtype=float64),
  params=Arap__Params(
         	mu=array(shape=(4802,), dtype=float64),
         )
)
forward    = Forward(
  model=Model(
    dirichlet=Dirichlet(
      dim=3,
      dirichlet_index=i64[534](jax),
      dirichlet_val...tart_threshold=Array(inf, dtype=float64, weak_type=True),
    max_delta=Array(inf, dtype=float64, weak_type=True)
  )
)
mesh       = UnstructuredGrid (0x7f3bdd1c27a0)
  N Cells:    4802
  N Points:   1317
  X Bounds:   3.159e+00, 3.892e+00
  Y Bounds:   -8.999e-02, 7.415e-01
  Z Bounds:   -1.735e-18, 3.000e-01
  N Arrays:   4
model      = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=i64[534](jax),
    dirichlet_value=f64[534](jax),
    fre... dtype=float64),
                 )
        )
      }
    )
  ),
  edges_length_mean=Array(0.05821938, dtype=float64)
)
rng        = Generator(PCG64) at 0x7F3BDC7774C0
solution   = OptimizeSolution(
  result=<Result.STAGNATION: 'stagnation'>,
  state=PNCGState(
    structure=Structure(full_flat=f64...PNCGStats(
    n_steps=372,
    relative_decrease=Array(1.9753192e-11, dtype=float64),
    time=50.61154793199967
  )
)

#x1B[1m#x1B[31mtests/model/test_forward.py#x1B[0m:32: AssertionError

tests/model/test_inverse.py::test_inverse

Flake rate in main: 97.44% (Passed 1 times, Failed 38 times)

Stack Traces | 36s run time

>   #x1B[0msolution: Optimizer.Solution = inverse.solve(params)#x1B[90m#x1B[39;49;00m


#x1B[1m#x1B[31mtests/model/test_inverse.py#x1B[0m:93: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
#x1B[1m#x1B[.../apple/inverse/_inverse.py#x1B[0m:189: in solve
    #x1B[0moptimizer_solution: Optimizer.Solution = #x1B[96mself#x1B[39;49;00m.optimizer.minimize(#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m.venv/lib/python3.12.../optim/scipy/_scipy.py#x1B[0m:117: in minimize
    #x1B[0mraw: OptimizeResult = scipy.optimize.minimize(  #x1B[90m# pyright: ignore[reportCallIssue]#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m.venv/lib/python3.12.../scipy/optimize/_minimize.py#x1B[0m:784: in minimize
    #x1B[0mres = _minimize_lbfgsb(fun, x0, args, jac, bounds,#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m.venv/lib/python3.12.../scipy/optimize/_lbfgsb_py.py#x1B[0m:413: in _minimize_lbfgsb
    #x1B[0msf = _prepare_scalar_function(fun, x0, jac=jac, args=args, epsilon=eps,#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m.venv/lib/python3.12.../scipy/optimize/_optimize.py#x1B[0m:310: in _prepare_scalar_function
    #x1B[0msf = ScalarFunction(fun, x0, args, grad, hess,#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m.venv/lib/python3.12.../scipy/optimize/_differentiable_functions.py#x1B[0m:283: in __init__
    #x1B[0m#x1B[96mself#x1B[39;49;00m._update_fun()#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m.venv/lib/python3.12.../scipy/optimize/_differentiable_functions.py#x1B[0m:362: in _update_fun
    #x1B[0mfx = #x1B[96mself#x1B[39;49;00m._wrapped_fun(#x1B[96mself#x1B[39;49;00m.x)#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m.venv/lib/python3.12.../scipy/_lib/_util.py#x1B[0m:603: in __call__
    #x1B[0mfx = #x1B[96mself#x1B[39;49;00m.f(np.copy(x), *#x1B[96mself#x1B[39;49;00m.args)#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m.venv/lib/python3.12.../scipy/optimize/_optimize.py#x1B[0m:80: in __call__
    #x1B[0m#x1B[96mself#x1B[39;49;00m._compute_if_needed(x, *args)#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m.venv/lib/python3.12.../scipy/optimize/_optimize.py#x1B[0m:74: in _compute_if_needed
    #x1B[0mfg = #x1B[96mself#x1B[39;49;00m.fun(x, *args)#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m.venv/lib/python3.12.../peach/functools/_descriptor.py#x1B[0m:62: in wrapper
    #x1B[0moutputs: Sequence[Any] = _as_tuple(wrapped(*args, **kwargs))#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[.../apple/inverse/_inverse.py#x1B[0m:206: in value_and_grad
    #x1B[0mp: Full = #x1B[96mself#x1B[39;49;00m.adjoint(u_full, dLdu)#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[.../apple/inverse/_inverse.py#x1B[0m:98: in adjoint
    #x1B[0msolution: LinearSolver.Solution = #x1B[96mself#x1B[39;49;00m.adjoint_inner(u, dLdu)#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[.../apple/inverse/_inverse.py#x1B[0m:138: in adjoint_inner
    #x1B[0msolution: LinearSolver.Solution = #x1B[96mself#x1B[39;49;00m.adjoint_solver.solve(system, params)#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m.venv/lib/python3.12.../linalg/abc/_solver.py#x1B[0m:76: in solve
    #x1B[0mstate, stats, result = #x1B[96mself#x1B[39;49;00m._solve(#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m.venv/lib/python3.12.../linalg/misc/_composite.py#x1B[0m:57: in _solve
    #x1B[0msolution = solver.solve(#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m.venv/lib/python3.12.../linalg/abc/_solver.py#x1B[0m:76: in solve
    #x1B[0mstate, stats, result = #x1B[96mself#x1B[39;49;00m._solve(#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m.venv/lib/python3.12.../linalg/jax/_base.py#x1B[0m:70: in _solve
    #x1B[0mstate.params_flat, stats.info = #x1B[96mself#x1B[39;49;00m._wrapped(#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m.venv/lib/python3.12.../linalg/jax/_cg.py#x1B[0m:17: in _wrapped
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m jax.scipy.sparse.linalg.cg(*args, **kwargs)#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m.venv/lib/python3.12.../scipy/sparse/linalg.py#x1B[0m:286: in cg
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m _isolve(_cg_solve,#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m.venv/lib/python3.12.../scipy/sparse/linalg.py#x1B[0m:226: in _isolve
    #x1B[0mx = lax.custom_linear_solve(#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31m.venv/lib/python3.12.../scipy/sparse/linalg.py#x1B[0m:128: in _cg_solve
    #x1B[0mr0 = _sub(b, A(x0))#x1B[90m#x1B[39;49;00m
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

>   #x1B[0m#x1B[94mreturn#x1B[39;49;00m wrapped(*args, **kwargs)#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31mE   jax._src.source_info_util.JaxStackTraceBeforeTransformation: jax.errors.JaxRuntimeError: NOT_FOUND: No FFI handler registered for WarpModelAdapter___hess_prod_callable__locals__hess_prod_callable_825 on a platform Host (canonical host)#x1B[0m
#x1B[1m#x1B[31mE   #x1B[0m
#x1B[1m#x1B[31mE   The preceding stack trace is the source of the JAX operation that, once transformed by JAX, triggered the following exception.#x1B[0m
#x1B[1m#x1B[31mE   #x1B[0m
#x1B[1m#x1B[31mE   --------------------#x1B[0m


#x1B[1m#x1B[31m.venv/lib/python3.12.../grapes/timing/_callable.py#x1B[0m:26: JaxStackTraceBeforeTransformation

#x1B[33mThe above exception was the direct cause of the following exception:#x1B[0m

mesh = UnstructuredGrid (0x7f3bdd5b1d80)
  N Cells:    96
  N Points:   35
  X Bounds:   0.000e+00, 2.000e+00
  Y Bounds:   0.000e+00, 2.000e+00
  Z Bounds:   0.000e+00, 2.000e+00
  N Arrays:   9
model = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=i64[27](jax),
    dirichlet_value=f64[27](jax),
    free_...,
                 )
        )
      }
    )
  ),
  edges_length_mean=Array(1.04115139, dtype=float64),
  frozen=True
)

    #x1B[0m#x1B[94mdef#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[92mtest_inverse#x1B[39;49;00m(mesh: pv.UnstructuredGrid, model: Model) -> #x1B[94mNone#x1B[39;49;00m:#x1B[90m#x1B[39;49;00m
        #x1B[37m@tree#x1B[39;49;00m.define#x1B[90m#x1B[39;49;00m
        #x1B[94mclass#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[04m#x1B[92mInverseActivation#x1B[39;49;00m(Inverse):#x1B[90m#x1B[39;49;00m
            surface_idx: Integer[Array, #x1B[33m"#x1B[39;49;00m#x1B[33m surface#x1B[39;49;00m#x1B[33m"#x1B[39;49;00m]#x1B[90m#x1B[39;49;00m
            target: Float[Array, #x1B[33m"#x1B[39;49;00m#x1B[33msurface dim#x1B[39;49;00m#x1B[33m"#x1B[39;49;00m]#x1B[90m#x1B[39;49;00m
    #x1B[90m#x1B[39;49;00m
            #x1B[37m@tree#x1B[39;49;00m.define#x1B[90m#x1B[39;49;00m
            #x1B[94mclass#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[04m#x1B[92mParams#x1B[39;49;00m(Inverse.Params):#x1B[90m#x1B[39;49;00m
                activation: Float[Array, #x1B[33m"#x1B[39;49;00m#x1B[33mcells 6#x1B[39;49;00m#x1B[33m"#x1B[39;49;00m]#x1B[90m#x1B[39;49;00m
    #x1B[90m#x1B[39;49;00m
            #x1B[37m@tree#x1B[39;49;00m.define#x1B[90m#x1B[39;49;00m
            #x1B[94mclass#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[04m#x1B[92mAux#x1B[39;49;00m(Inverse.Aux):#x1B[90m#x1B[39;49;00m
                #x1B[94mpass#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
    #x1B[90m#x1B[39;49;00m
            #x1B[37m@override#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
            #x1B[94mdef#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[92mloss#x1B[39;49;00m(#x1B[90m#x1B[39;49;00m
                #x1B[96mself#x1B[39;49;00m, u: Float[Array, #x1B[33m"#x1B[39;49;00m#x1B[33mpoints dim#x1B[39;49;00m#x1B[33m"#x1B[39;49;00m], params: ModelParams#x1B[90m#x1B[39;49;00m
            ) -> #x1B[96mtuple#x1B[39;49;00m[Float[Array, #x1B[33m"#x1B[39;49;00m#x1B[33m"#x1B[39;49;00m], Aux]:#x1B[90m#x1B[39;49;00m
                loss: Float[Array, #x1B[33m"#x1B[39;49;00m#x1B[33m"#x1B[39;49;00m] = #x1B[94m0.5#x1B[39;49;00m * jnp.sum(#x1B[90m#x1B[39;49;00m
                    jnp.square(u[#x1B[96mself#x1B[39;49;00m.surface_idx] - #x1B[96mself#x1B[39;49;00m.target)#x1B[90m#x1B[39;49;00m
                )#x1B[90m#x1B[39;49;00m
                #x1B[94mreturn#x1B[39;49;00m loss, #x1B[96mself#x1B[39;49;00m.Aux()#x1B[90m#x1B[39;49;00m
    #x1B[90m#x1B[39;49;00m
            #x1B[37m@override#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
            #x1B[94mdef#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[92mmake_params#x1B[39;49;00m(#x1B[96mself#x1B[39;49;00m, params: Params) -> ModelParams:  #x1B[90m# pyright: ignore[reportIncompatibleMethodOverride]#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
                #x1B[94mreturn#x1B[39;49;00m {#x1B[33m"#x1B[39;49;00m#x1B[33melastic#x1B[39;49;00m#x1B[33m"#x1B[39;49;00m: {#x1B[33m"#x1B[39;49;00m#x1B[33mactivation#x1B[39;49;00m#x1B[33m"#x1B[39;49;00m: params.activation}}#x1B[90m#x1B[39;49;00m
    #x1B[90m#x1B[39;49;00m
        forward = Forward(model)#x1B[90m#x1B[39;49;00m
        forward.step()#x1B[90m#x1B[39;49;00m
    #x1B[90m#x1B[39;49;00m
        surface: pv.PolyData = mesh.extract_surface()  #x1B[90m# pyright: ignore[reportAssignmentType]#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
        surface_idx: Integer[Array, #x1B[33m"#x1B[39;49;00m#x1B[33m surface#x1B[39;49;00m#x1B[33m"#x1B[39;49;00m] = jnp.asarray(surface.point_data[POINT_ID])#x1B[90m#x1B[39;49;00m
        target: Float[Array, #x1B[33m"#x1B[39;49;00m#x1B[33msurface dim#x1B[39;49;00m#x1B[33m"#x1B[39;49;00m] = forward.u_full[surface_idx]#x1B[90m#x1B[39;49;00m
    #x1B[90m#x1B[39;49;00m
        inverse = InverseActivation(forward, surface_idx=surface_idx, target=target)#x1B[90m#x1B[39;49;00m
        params = InverseActivation.Params(activation=jnp.zeros((mesh.n_cells, #x1B[94m6#x1B[39;49;00m)))#x1B[90m#x1B[39;49;00m
>       solution: Optimizer.Solution = inverse.solve(params)#x1B[90m#x1B[39;49;00m
                                       ^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m

InverseActivation = <class 'tests.model.test_inverse.test_inverse.<locals>.InverseActivation'>
forward    = Forward(
  model=Model(
    dirichlet=Dirichlet(
      dim=3,
      dirichlet_index=i64[27](jax),
      dirichlet_valu... beta_restart_threshold=Array(10., dtype=float64, weak_type=True),
    max_delta=Array(0.15617271, dtype=float64)
  )
)
inverse    = test_inverse.<locals>.InverseActivation(
  forward=Forward(
    model=Model(
      dirichlet=Dirichlet(
        dim=3,...lse, dtype=bool),
  last_forward_success=Array(True, dtype=bool),
  surface_idx=i64[26](jax),
  target=f64[26,3](jax)
)
mesh       = UnstructuredGrid (0x7f3bdd5b1d80)
  N Cells:    96
  N Points:   35
  X Bounds:   0.000e+00, 2.000e+00
  Y Bounds:   0.000e+00, 2.000e+00
  Z Bounds:   0.000e+00, 2.000e+00
  N Arrays:   9
model      = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=i64[27](jax),
    dirichlet_value=f64[27](jax),
    free_...,
                 )
        )
      }
    )
  ),
  edges_length_mean=Array(1.04115139, dtype=float64),
  frozen=True
)
params     = test_inverse.<locals>.InverseActivation.Params(activation=f64[96,6](jax))
surface    = PolyData (0x7f3bdcf37f40)
  N Cells:    48
  N Points:   26
  N Strips:   0
  X Bounds:   0.000e+00, 2.000e+00
  Y Bounds:   0.000e+00, 2.000e+00
  Z Bounds:   0.000e+00, 2.000e+00
  N Arrays:   11
surface_idx = Array([ 0,  4,  1,  3, 10,  9, 12,  5,  2, 11, 14,  7,  6, 15,  8, 17, 16,
       19, 18, 21, 20, 23, 24, 26, 25, 22], dtype=int64)
target     = Array([[ 0.        ,  0.        ,  0.        ],
       [ 0.11017376, -0.09044607,  0.36066905],
       [ 0.        ,  ...6942 ],
       [ 0.03024186, -0.38203043,  0.06045359],
       [ 0.0644187 , -0.19728744,  0.11883095]], dtype=float64)

#x1B[1m#x1B[31mtests/model/test_inverse.py#x1B[0m:93: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
#x1B[1m#x1B[.../apple/inverse/_inverse.py#x1B[0m:189: in solve
    #x1B[0moptimizer_solution: Optimizer.Solution = #x1B[96mself#x1B[39;49;00m.optimizer.minimize(#x1B[90m#x1B[39;49;00m
        callback   = None
        constraints = ()
        objective  = Objective(
  _kwargs={},
  _grad_wrapped=<bound method Inverse.grad of test_inverse.<locals>.InverseActivation(
      ...             surface_idx=i64[26](jax),
                            target=f64[26,3](jax)
                          )>
)
        params     = test_inverse.<locals>.InverseActivation.Params(activation=f64[96,6](jax))
        self       = test_inverse.<locals>.InverseActivation(
  forward=Forward(
    model=Model(
      dirichlet=Dirichlet(
        dim=3,...lse, dtype=bool),
  last_forward_success=Array(True, dtype=bool),
  surface_idx=i64[26](jax),
  target=f64[26,3](jax)
)
#x1B[1m#x1B[31m.venv/lib/python3.12.../optim/scipy/_scipy.py#x1B[0m:117: in minimize
    #x1B[0mraw: OptimizeResult = scipy.optimize.minimize(  #x1B[90m# pyright: ignore[reportCallIssue]#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
        callback   = None
        callback_wrapper = <FunctionWrapper at 0x7f3bdcf566c0 for function at 0x7f3bdcf46e80>
        constraints = []
        fun        = <function FunctionDescriptor.__get__.<locals>.wrapper at 0x7f3bdcf46d40>
        jac        = True
        objective  = Objective(
  structure=Structure(
    full_flat=f64[576](jax),
    static=test_inverse.<locals>.InverseActivation.Para...](jax)
                          )>,
  _value_and_grad_wrapper=<function FunctionDescriptor.__get__.<locals>.wrapper>
)
        options    = {'maxiter': 256}
        params     = test_inverse.<locals>.InverseActivation.Params(activation=f64[96,6](jax))
        self       = ScipyOptimizer(method='L-BFGS-B', tol=1e-05)
        state      = ScipyState(
  structure=Structure(
    full_flat=f64[576](jax),
    static=test_inverse.<locals>.InverseActivation.Params(activation=None)
  ),
  result={'x': f64[576](jax)}
)
        stats      = ScipyStats(time=10.655189682999662)
#x1B[1m#x1B[31m.venv/lib/python3.12.../scipy/optimize/_minimize.py#x1B[0m:784: in minimize
    #x1B[0mres = _minimize_lbfgsb(fun, x0, args, jac, bounds,#x1B[90m#x1B[39;49;00m
        args       = ()
        bounds     = None
        callback   = <function _wrap_callback.<locals>.wrapped_callback at 0x7f3bdcf46b60>
        constraints = []
        fun        = <scipy.optimize._optimize.MemoizeJac object at 0x7f3bdd484320>
        hess       = None
        hessp      = None
        jac        = <bound method MemoizeJac.derivative of <scipy.optimize._optimize.MemoizeJac object at 0x7f3bdd484320>>
        meth       = 'l-bfgs-b'
        method     = 'L-BFGS-B'
        options    = {'ftol': 1e-05, 'gtol': 1e-05, 'maxiter': 256}
        remove_vars = False
        tol        = 1e-05
        x0         = array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., ...0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])
#x1B[1m#x1B[31m.venv/lib/python3.12.../scipy/optimize/_lbfgsb_py.py#x1B[0m:413: in _minimize_lbfgsb
    #x1B[0msf = _prepare_scalar_function(fun, x0, jac=jac, args=args, epsilon=eps,#x1B[90m#x1B[39;49;00m
        args       = ()
        bounds     = None
        callback   = <function _wrap_callback.<locals>.wrapped_callback at 0x7f3bdcf46b60>
        disp       = <object object at 0x7f3c92313160>
        eps        = 1e-08
        factr      = np.float64(45035996273.70496)
        finite_diff_rel_step = None
        ftol       = 1e-05
        fun        = <scipy.optimize._optimize.MemoizeJac object at 0x7f3bdd484320>
        gtol       = 1e-05
        iprint     = <object object at 0x7f3c92313160>
        jac        = <bound method MemoizeJac.derivative of <scipy.optimize._optimize.MemoizeJac object at 0x7f3bdd484320>>
        m          = 10
        maxcor     = 10
        maxfun     = 15000
        maxiter    = 256
        maxls      = 20
        n          = 576
        pgtol      = 1e-05
        unknown_options = {}
        workers    = None
        x0         = array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., ...0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])
#x1B[1m#x1B[31m.venv/lib/python3.12.../scipy/optimize/_optimize.py#x1B[0m:310: in _prepare_scalar_function
    #x1B[0msf = ScalarFunction(fun, x0, args, grad, hess,#x1B[90m#x1B[39;49;00m
        args       = ()
        bounds     = (-inf, inf)
        epsilon    = 1e-08
        finite_diff_rel_step = None
        fun        = <scipy.optimize._optimize.MemoizeJac object at 0x7f3bdd484320>
        grad       = <bound method MemoizeJac.derivative of <scipy.optimize._optimize.MemoizeJac object at 0x7f3bdd484320>>
        hess       = <function _prepare_scalar_function.<locals>.hess at 0x7f3bdcf47420>
        jac        = <bound method MemoizeJac.derivative of <scipy.optimize._optimize.MemoizeJac object at 0x7f3bdd484320>>
        workers    = <class 'map'>
        x0         = array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., ...0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])
#x1B[1m#x1B[31m.venv/lib/python3.12.../scipy/optimize/_differentiable_functions.py#x1B[0m:283: in __init__
    #x1B[0m#x1B[96mself#x1B[39;49;00m._update_fun()#x1B[90m#x1B[39;49;00m
        _dtype     = dtype('float64')
        _x         = array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., ...0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])
        args       = ()
        epsilon    = 1e-08
        finite_diff_bounds = (-inf, inf)
        finite_diff_options = {}
        finite_diff_rel_step = None
        fun        = <scipy.optimize._optimize.MemoizeJac object at 0x7f3bdd484320>
        grad       = <bound method MemoizeJac.derivative of <scipy.optimize._optimize.MemoizeJac object at 0x7f3bdd484320>>
        hess       = <function _prepare_scalar_function.<locals>.hess at 0x7f3bdcf47420>
        self       = <scipy.optimize._differentiable_functions.ScalarFunction object at 0x7f3bdd487950>
        workers    = <class 'map'>
        x0         = array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., ...0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])
        xp         = <module 'scipy._lib.array_api_compat.numpy' from '.../apple/apple/.venv/lib/python3.12.../array_api_compat/numpy/__init__.py'>
#x1B[1m#x1B[31m.venv/lib/python3.12.../scipy/optimize/_differentiable_functions.py#x1B[0m:362: in _update_fun
    #x1B[0mfx = #x1B[96mself#x1B[39;49;00m._wrapped_fun(#x1B[96mself#x1B[39;49;00m.x)#x1B[90m#x1B[39;49;00m
         ^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        self       = <scipy.optimize._differentiable_functions.ScalarFunction object at 0x7f3bdd487950>
#x1B[1m#x1B[31m.venv/lib/python3.12.../scipy/_lib/_util.py#x1B[0m:603: in __call__
    #x1B[0mfx = #x1B[96mself#x1B[39;49;00m.f(np.copy(x), *#x1B[96mself#x1B[39;49;00m.args)#x1B[90m#x1B[39;49;00m
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        self       = <scipy._lib._util._ScalarFunctionWrapper object at 0x7f3bdd48ee70>
        x          = array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., ...0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])
#x1B[1m#x1B[31m.venv/lib/python3.12.../scipy/optimize/_optimize.py#x1B[0m:80: in __call__
    #x1B[0m#x1B[96mself#x1B[39;49;00m._compute_if_needed(x, *args)#x1B[90m#x1B[39;49;00m
        args       = ()
        self       = <scipy.optimize._optimize.MemoizeJac object at 0x7f3bdd484320>
        x          = array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., ...0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])
#x1B[1m#x1B[31m.venv/lib/python3.12.../scipy/optimize/_optimize.py#x1B[0m:74: in _compute_if_needed
    #x1B[0mfg = #x1B[96mself#x1B[39;49;00m.fun(x, *args)#x1B[90m#x1B[39;49;00m
         ^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = ()
        self       = <scipy.optimize._optimize.MemoizeJac object at 0x7f3bdd484320>
        x          = array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., ...0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.])
#x1B[1m#x1B[.../apple/inverse/_inverse.py#x1B[0m:206: in value_and_grad
    #x1B[0mp: Full = #x1B[96mself#x1B[39;49;00m.adjoint(u_full, dLdu)#x1B[90m#x1B[39;49;00m
              ^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        aux        = test_inverse.<locals>.InverseActivation.Aux()
        dLdq       = {'elastic': {'activation': Array([[0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0...],
       [0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0.]], dtype=float64)}}
        dLdu       = Array([[ 0.        ,  0.        ,  0.        ],
       [ 0.        ,  0.        ,  0.        ],
       [ 0.        ,  ...     ],
       [ 0.        ,  0.        ,  0.        ],
       [ 0.        ,  0.        ,  0.        ]], dtype=float64)
        loss       = Array(1.58222216, dtype=float64)
        model_params = {'elastic': {'activation': Array([[0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0...],
       [0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0.],
       [0., 0., 0., 0., 0., 0.]], dtype=float64)}}
        model_params_vjp = VJP(fun=functools.partial(<function _vjp3_callable at 0x7f3c932ba200>, [], [False], { lambda ; a:f64[96,6]. let  in (a...arams[()], [*]),)), out_tree=PyTreeDef({'elastic': {'activation': *}}), args_res=[(NotNeeded(),)], opaque_residuals=[])
        params     = test_inverse.<locals>.InverseActivation.Params(activation=f64[96,6](jax))
        self       = test_inverse.<locals>.InverseActivation(
  forward=Forward(
    model=Model(
      dirichlet=Dirichlet(
        dim=3,...lse, dtype=bool),
  last_forward_success=Array(True, dtype=bool),
  surface_idx=i64[26](jax),
  target=f64[26,3](jax)
)
        u_full     = Array([[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00],
       [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00]...00,  0.00000000e+00,  0.00000000e+00],
       [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00]],      dtype=float64)
#x1B[1m#x1B[.../apple/inverse/_inverse.py#x1B[0m:98: in adjoint
    #x1B[0msolution: LinearSolver.Solution = #x1B[96mself#x1B[39;49;00m.adjoint_inner(u, dLdu)#x1B[90m#x1B[39;49;00m
                                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        dLdu       = Array([[ 0.        ,  0.        ,  0.        ],
       [ 0.        ,  0.        ,  0.        ],
       [ 0.        ,  ...     ],
       [ 0.        ,  0.        ,  0.        ],
       [ 0.        ,  0.        ,  0.        ]], dtype=float64)
        self       = test_inverse.<locals>.InverseActivation(
  forward=Forward(
    model=Model(
      dirichlet=Dirichlet(
        dim=3,...lse, dtype=bool),
  last_forward_success=Array(True, dtype=bool),
  surface_idx=i64[26](jax),
  target=f64[26,3](jax)
)
        u          = Array([[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00],
       [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00]...00,  0.00000000e+00,  0.00000000e+00],
       [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00]],      dtype=float64)
#x1B[1m#x1B[.../apple/inverse/_inverse.py#x1B[0m:138: in adjoint_inner
    #x1B[0msolution: LinearSolver.Solution = #x1B[96mself#x1B[39;49;00m.adjoint_solver.solve(system, params)#x1B[90m#x1B[39;49;00m
                                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        dLdu       = Array([[ 0.        ,  0.        ,  0.        ],
       [ 0.        ,  0.        ,  0.        ],
       [ 0.        ,  ...     ],
       [ 0.        ,  0.        ,  0.        ],
       [ 0.        ,  0.        ,  0.        ]], dtype=float64)
        matvec     = <jax._src.custom_derivatives.custom_jvp object at 0x7f3bde33cec0>
        matvec_jvp = <function Inverse.adjoint_inner.<locals>.matvec_jvp at 0x7f3bdd1ac4a0>
        params     = Array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., ... 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float64)
        preconditioner = Array([0.31578947, 0.27272727, 0.31578947, 0.16666667, 0.16666667,
       0.16666667, 0.35294118, 0.42857143, 0.352941... 0.125     , 0.125     , 0.125     , 0.125     , 0.125     ,
       0.125     , 0.125     , 0.125     ], dtype=float64)
        preconditioner_fn = <function Inverse.adjoint_inner.<locals>.preconditioner_fn at 0x7f3bdd1ad120>
        self       = test_inverse.<locals>.InverseActivation(
  forward=Forward(
    model=Model(
      dirichlet=Dirichlet(
        dim=3,...lse, dtype=bool),
  last_forward_success=Array(True, dtype=bool),
  surface_idx=i64[26](jax),
  target=f64[26,3](jax)
)
        system     = LinearSystem(
  _kwargs={},
  _matvec_wrapped=<jax._src.custom_derivatives.custom_jvp object at 0x7f3bde33cec0>,
  b=f...r.<locals>.preconditioner_fn>,
  _rpreconditioner_wrapped=<function Inverse.adjoint_inner.<locals>.preconditioner_fn>
)
        u          = Array([[ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00],
       [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00]...00,  0.00000000e+00,  0.00000000e+00],
       [ 0.00000000e+00,  0.00000000e+00,  0.00000000e+00]],      dtype=float64)
        u_free     = Array([ 5.75383323e-17,  1.52721599e-17,  2.58720539e-17,  6.07349063e-17,
        2.60756920e-17,  6.41272332e-17,  6...000000e+00,  0.00000000e+00,  0.00000000e+00,  0.00000000e+00,
        0.00000000e+00,  0.00000000e+00], dtype=float64)
#x1B[1m#x1B[31m.venv/lib/python3.12.../linalg/abc/_solver.py#x1B[0m:76: in solve
    #x1B[0mstate, stats, result = #x1B[96mself#x1B[39;49;00m._solve(#x1B[90m#x1B[39;49;00m
        StateT     = StateT
        StatsT     = StatsT
        callback   = None
        constraints = []
        params     = Array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., ... 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float64)
        self       = CompositeSolver(
  jit=True,
  timer=True,
  solvers=[
    JaxCG(
      max_steps=1000,
      atol=Array(0., dtype=flo...  continue_atol=Array(0., dtype=float64, weak_type=True),
  continue_rtol=Array(0.001, dtype=float64, weak_type=True)
)
        state      = CompositeState(
  structure=Structure(full_flat=f64[78](jax)), params_flat=f64[78](jax), state=[]
)
        stats      = CompositeStats(stats=[], time=0.805306123000264)
        system     = LinearSystem(
  structure=Structure(full_flat=f64[78](jax)),
  _flatten=True,
  _jit=True,
  _kwargs={},
  _timer=True...r.<locals>.preconditioner_fn>,
  _rpreconditioner_wrapped=<function Inverse.adjoint_inner.<locals>.preconditioner_fn>
)
#x1B[1m#x1B[31m.venv/lib/python3.12.../linalg/misc/_composite.py#x1B[0m:57: in _solve
    #x1B[0msolution = solver.solve(#x1B[90m#x1B[39;49;00m
        callback   = None
        constraints = []
        self       = CompositeSolver(
  jit=True,
  timer=True,
  solvers=[
    JaxCG(
      max_steps=1000,
      atol=Array(0., dtype=flo...  continue_atol=Array(0., dtype=float64, weak_type=True),
  continue_rtol=Array(0.001, dtype=float64, weak_type=True)
)
        solution   = None
        solver     = JaxCG(
  max_steps=1000,
  atol=Array(0., dtype=float64, weak_type=True),
  rtol=Array(0.001, dtype=float64, weak_type...
  atol_primary=Array(0., dtype=float64, weak_type=True),
  rtol_primary=Array(1.e-05, dtype=float64, weak_type=True)
)
        state      = CompositeState(
  structure=Structure(full_flat=f64[78](jax)), params_flat=f64[78](jax), state=[]
)
        stats      = CompositeStats(stats=[], time=0.8168954270004178)
        system     = LinearSystem(
  structure=Structure(full_flat=f64[78](jax)),
  _flatten=True,
  _jit=True,
  _kwargs={},
  _timer=True...r.<locals>.preconditioner_fn>,
  _rpreconditioner_wrapped=<function Inverse.adjoint_inner.<locals>.preconditioner_fn>
)
#x1B[1m#x1B[31m.venv/lib/python3.12.../linalg/abc/_solver.py#x1B[0m:76: in solve
    #x1B[0mstate, stats, result = #x1B[96mself#x1B[39;49;00m._solve(#x1B[90m#x1B[39;49;00m
        StateT     = StateT
        StatsT     = StatsT
        callback   = None
        constraints = []
        params     = Array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., ... 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float64)
        self       = JaxCG(
  max_steps=1000,
  atol=Array(0., dtype=float64, weak_type=True),
  rtol=Array(0.001, dtype=float64, weak_type...
  atol_primary=Array(0., dtype=float64, weak_type=True),
  rtol_primary=Array(1.e-05, dtype=float64, weak_type=True)
)
        state      = JaxState(structure=Structure(full_flat=f64[78](jax)), params_flat=f64[78](jax))
        stats      = JaxStats(time=0.8223094720005975)
        system     = LinearSystem(
  structure=Structure(full_flat=f64[78](jax)),
  _flatten=True,
  _jit=True,
  _kwargs={},
  _timer=True...JitWrapper at 0x7f3bdd1f1b20>,
  _rpreconditioner_wrapped=<function Inverse.adjoint_inner.<locals>.preconditioner_fn>
)
#x1B[1m#x1B[31m.venv/lib/python3.12.../linalg/jax/_base.py#x1B[0m:70: in _solve
    #x1B[0mstate.params_flat, stats.info = #x1B[96mself#x1B[39;49;00m._wrapped(#x1B[90m#x1B[39;49;00m
        callback   = None
        constraints = []
        self       = JaxCG(
  max_steps=1000,
  atol=Array(0., dtype=float64, weak_type=True),
  rtol=Array(0.001, dtype=float64, weak_type...
  atol_primary=Array(0., dtype=float64, weak_type=True),
  rtol_primary=Array(1.e-05, dtype=float64, weak_type=True)
)
        state      = JaxState(structure=Structure(full_flat=f64[78](jax)), params_flat=f64[78](jax))
        stats      = JaxStats(time=0.8296433830000751)
        system     = LinearSystem(
  structure=Structure(full_flat=f64[78](jax)),
  _flatten=True,
  _jit=True,
  _kwargs={},
  _timer=True...JitWrapper at 0x7f3bdd1f1b20>,
  _rpreconditioner_wrapped=<function Inverse.adjoint_inner.<locals>.preconditioner_fn>
)
#x1B[1m#x1B[31m.venv/lib/python3.12.../linalg/jax/_cg.py#x1B[0m:17: in _wrapped
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m jax.scipy.sparse.linalg.cg(*args, **kwargs)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = (<FunctionWrapper at 0x7f3bdd294120 for _JitWrapper at 0x7f3bdceaa1b0>, Array([ 0.14731738, -0.14202482,  0.35077357, ...0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float64))
        kwargs     = {'M': <FunctionWrapper at 0x7f3bdd296030 for _JitWrapper at 0x7f3bdd1f1b20>, 'atol': Array(0., dtype=float64, weak_type=True), 'maxiter': 1000, 'tol': Array(1.e-05, dtype=float64, weak_type=True)}
        self       = JaxCG(
  max_steps=1000,
  atol=Array(0., dtype=float64, weak_type=True),
  rtol=Array(0.001, dtype=float64, weak_type...
  atol_primary=Array(0., dtype=float64, weak_type=True),
  rtol_primary=Array(1.e-05, dtype=float64, weak_type=True)
)
#x1B[1m#x1B[31m.venv/lib/python3.12.../scipy/sparse/linalg.py#x1B[0m:286: in cg
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m _isolve(_cg_solve,#x1B[90m#x1B[39;49;00m
        A          = <FunctionWrapper at 0x7f3bdd294120 for _JitWrapper at 0x7f3bdceaa1b0>
        M          = <FunctionWrapper at 0x7f3bdd296030 for _JitWrapper at 0x7f3bdd1f1b20>
        atol       = Array(0., dtype=float64, weak_type=True)
        b          = Array([ 0.14731738, -0.14202482,  0.35077357,  0.11017376, -0.09044607,
        0.36066905, -0.03154436, -0.06405633, ...    , -0.        , -0.        , -0.        , -0.        ,
       -0.        , -0.        , -0.        ], dtype=float64)
        maxiter    = 1000
        tol        = Array(1.e-05, dtype=float64, weak_type=True)
        x0         = Array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., ... 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float64)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

_isolve_solve = <function _cg_solve at 0x7f3bdd1aeb60>
A = <FunctionWrapper at 0x7f3bdd294120 for _JitWrapper at 0x7f3bdceaa1b0>
b = Array([ 0.14731738, -0.14202482,  0.35077357,  0.11017376, -0.09044607,
        0.36066905, -0.03154436, -0.06405633, ...    , -0.        , -0.        , -0.        , -0.        ,
       -0.        , -0.        , -0.        ], dtype=float64)
x0 = Array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., ... 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float64)

    #x1B[0m#x1B[94mdef#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[92m_isolve#x1B[39;49;00m(_isolve_solve, A, b, x0=#x1B[94mNone#x1B[39;49;00m, *, tol=#x1B[94m1e-5#x1B[39;49;00m, atol=#x1B[94m0.0#x1B[39;49;00m,#x1B[90m#x1B[39;49;00m
                maxiter=#x1B[94mNone#x1B[39;49;00m, M=#x1B[94mNone#x1B[39;49;00m, check_symmetric=#x1B[94mFalse#x1B[39;49;00m):#x1B[90m#x1B[39;49;00m
      #x1B[94mif#x1B[39;49;00m x0 #x1B[95mis#x1B[39;49;00m #x1B[94mNone#x1B[39;49;00m:#x1B[90m#x1B[39;49;00m
        x0 = tree_map(jnp.zeros_like, b)#x1B[90m#x1B[39;49;00m
    #x1B[90m#x1B[39;49;00m
      b, x0 = api.device_put((b, x0))#x1B[90m#x1B[39;49;00m
    #x1B[90m#x1B[39;49;00m
      #x1B[94mif#x1B[39;49;00m maxiter #x1B[95mis#x1B[39;49;00m #x1B[94mNone#x1B[39;49;00m:#x1B[90m#x1B[39;49;00m
        size = #x1B[96msum#x1B[39;49;00m(bi.size #x1B[94mfor#x1B[39;49;00m bi #x1B[95min#x1B[39;49;00m tree_leaves(b))#x1B[90m#x1B[39;49;00m
        maxiter = #x1B[94m10#x1B[39;49;00m * size  #x1B[90m# copied from scipy#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
    #x1B[90m#x1B[39;49;00m
      #x1B[94mif#x1B[39;49;00m M #x1B[95mis#x1B[39;49;00m #x1B[94mNone#x1B[39;49;00m:#x1B[90m#x1B[39;49;00m
        M = _identity#x1B[90m#x1B[39;49;00m
      A = _normalize_matvec(A)#x1B[90m#x1B[39;49;00m
      M = _normalize_matvec(M)#x1B[90m#x1B[39;49;00m
    #x1B[90m#x1B[39;49;00m
      #x1B[94mif#x1B[39;49;00m tree_structure(x0) != tree_structure(b):#x1B[90m#x1B[39;49;00m
        #x1B[94mraise#x1B[39;49;00m #x1B[96mValueError#x1B[39;49;00m(#x1B[90m#x1B[39;49;00m
            #x1B[33m'#x1B[39;49;00m#x1B[33mx0 and b must have matching tree structure: #x1B[39;49;00m#x1B[33m'#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
            #x1B[33mf#x1B[39;49;00m#x1B[33m'#x1B[39;49;00m#x1B[33m{#x1B[39;49;00mtree_structure(x0)#x1B[33m}#x1B[39;49;00m#x1B[33m vs #x1B[39;49;00m#x1B[33m{#x1B[39;49;00mtree_structure(b)#x1B[33m}#x1B[39;49;00m#x1B[33m'#x1B[39;49;00m)#x1B[90m#x1B[39;49;00m
    #x1B[90m#x1B[39;49;00m
      #x1B[94mif#x1B[39;49;00m _shapes(x0) != _shapes(b):#x1B[90m#x1B[39;49;00m
        #x1B[94mraise#x1B[39;49;00m #x1B[96mValueError#x1B[39;49;00m(#x1B[90m#x1B[39;49;00m
            #x1B[33m'#x1B[39;49;00m#x1B[33marrays in x0 and b must have matching shapes: #x1B[39;49;00m#x1B[33m'#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
            #x1B[33mf#x1B[39;49;00m#x1B[33m'#x1B[39;49;00m#x1B[33m{#x1B[39;49;00m_shapes(x0)#x1B[33m}#x1B[39;49;00m#x1B[33m vs #x1B[39;49;00m#x1B[33m{#x1B[39;49;00m_shapes(b)#x1B[33m}#x1B[39;49;00m#x1B[33m'#x1B[39;49;00m)#x1B[90m#x1B[39;49;00m
    #x1B[90m#x1B[39;49;00m
      isolve_solve = partial(#x1B[90m#x1B[39;49;00m
          _isolve_solve, x0=x0, tol=tol, atol=atol, maxiter=maxiter, M=M)#x1B[90m#x1B[39;49;00m
    #x1B[90m#x1B[39;49;00m
      #x1B[90m# real-valued positive-definite linear operators are symmetric#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
      #x1B[94mdef#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[92mreal_valued#x1B[39;49;00m(x):#x1B[90m#x1B[39;49;00m
        #x1B[94mreturn#x1B[39;49;00m #x1B[95mnot#x1B[39;49;00m #x1B[96missubclass#x1B[39;49;00m(x.dtype.type, np.complexfloating)#x1B[90m#x1B[39;49;00m
      symmetric = #x1B[96mall#x1B[39;49;00m(#x1B[96mmap#x1B[39;49;00m(real_valued, tree_leaves(b))) \
        #x1B[94mif#x1B[39;49;00m check_symmetric #x1B[94melse#x1B[39;49;00m #x1B[94mFalse#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
>     x = lax.custom_linear_solve(#x1B[90m#x1B[39;49;00m
          A, b, solve=isolve_solve, transpose_solve=isolve_solve,#x1B[90m#x1B[39;49;00m
          symmetric=symmetric)#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31mE     jax.errors.JaxRuntimeError: NOT_FOUND: No FFI handler registered for WarpModelAdapter___hess_prod_callable__locals__hess_prod_callable_825 on a platform Host (canonical host)#x1B[0m
#x1B[1m#x1B[31mE     --------------------#x1B[0m
#x1B[1m#x1B[31mE     For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.#x1B[0m

A          = <FunctionWrapper at 0x7f3bdd294120 for _JitWrapper at 0x7f3bdceaa1b0>
M          = <FunctionWrapper at 0x7f3bdd296030 for _JitWrapper at 0x7f3bdd1f1b20>
_isolve_solve = <function _cg_solve at 0x7f3bdd1aeb60>
atol       = Array(0., dtype=float64, weak_type=True)
b          = Array([ 0.14731738, -0.14202482,  0.35077357,  0.11017376, -0.09044607,
        0.36066905, -0.03154436, -0.06405633, ...    , -0.        , -0.        , -0.        , -0.        ,
       -0.        , -0.        , -0.        ], dtype=float64)
check_symmetric = True
isolve_solve = functools.partial(<function _cg_solve at 0x7f3bdd1aeb60>, x0=Array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.... dtype=float64, weak_type=True), maxiter=1000, M=<FunctionWrapper at 0x7f3bdd296030 for _JitWrapper at 0x7f3bdd1f1b20>)
maxiter    = 1000
real_valued = <function _isolve.<locals>.real_valued at 0x7f3bdd1adc60>
symmetric  = True
tol        = Array(1.e-05, dtype=float64, weak_type=True)
x0         = Array([0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., ... 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.,
       0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], dtype=float64)

#x1B[1m#x1B[31m.venv/lib/python3.12.../scipy/sparse/linalg.py#x1B[0m:226: JaxRuntimeError

tests/warp/energies/elastic/hyperelastic/test_arap.py::test_arap_hess_prod

Flake rate in main: 66.67% (Passed 19 times, Failed 38 times)

Stack Traces | 59.9s run time

model = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
mesh = UnstructuredGrid (0x7f3bf8d15f00)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   2

    #x1B[0m#x1B[37m@hypothesis#x1B[39;49;00m.given(seed=testing.seed())#x1B[90m#x1B[39;49;00m
>   #x1B[94mdef#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[92mtest_arap_hess_prod#x1B[39;49;00m(seed: #x1B[96mint#x1B[39;49;00m, model: Model, mesh: pv.UnstructuredGrid) -> #x1B[94mNone#x1B[39;49;00m:#x1B[90m#x1B[39;49;00m
                   ^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m

f          = <function given.<locals>.run_test_as_given.<locals>.wrapped_test at 0x7f3c401e8360>
mesh       = UnstructuredGrid (0x7f3bf8d15f00)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   2
model      = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)

#x1B[1m#x1B[.../elastic/hyperelastic/test_arap.py#x1B[0m:59: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
#x1B[1m#x1B[.../elastic/hyperelastic/test_arap.py#x1B[0m:60: in test_arap_hess_prod
    #x1B[0mcommon.check_hess_prod(seed, model, mesh)#x1B[90m#x1B[39;49;00m
        mesh       = UnstructuredGrid (0x7f3bf8d15f00)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   2
        model      = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
        seed       = 0
#x1B[1m#x1B[.../elastic/hyperelastic/common.py#x1B[0m:55: in check_hess_prod
    #x1B[0mtesting.check_jvp(model.grad, model.hess_prod, u, rtol=#x1B[94m1e-3#x1B[39;49;00m)#x1B[90m#x1B[39;49;00m
        mesh       = UnstructuredGrid (0x7f3bf8d15f00)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   2
        model      = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
        seed       = 0
        u          = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
#x1B[1m#x1B[.../apple/model/_model.py#x1B[0m:117: in hess_prod
    #x1B[0moutput_wp: Full = #x1B[96mself#x1B[39;49;00m.warp.hess_prod(u_full, p_full)#x1B[90m#x1B[39;49;00m
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        output_jax = Array([[0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.]], dtype=float64)
        p          = Array([[-0.16308578, -0.56740909,  0.93064292],
       [ 0.14900107,  0.06445298, -0.29018964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64)
        p_full     = Array([[-0.16308578, -0.56740909,  0.93064292],
       [ 0.14900107,  0.06445298, -0.29018964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64)
        self       = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
        u          = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
        u_full     = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
#x1B[1m#x1B[.../warp/model/_adapter.py#x1B[0m:60: in hess_prod
    #x1B[0m(output,) = #x1B[96mself#x1B[39;49;00m._hess_prod_callable(u, p, output_dims=u.shape)#x1B[90m#x1B[39;49;00m
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        p          = Array([[-0.16308578, -0.56740909,  0.93064292],
       [ 0.14900107,  0.06445298, -0.29018964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64)
        self       = WarpModelAdapter(
  wrapped=WarpModel(
    energies={
      'elastic':
      Arap(
        id='elastic',
        requi...1,), dtype=float64),
               ),
        clamp_hess_diag=False,
        clamp_hess_quad=False
      )
    }
  )
)
        u          = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
#x1B[1m#x1B[31m.venv/lib/python3.12.../_src/jax_experimental/ffi.py#x1B[0m:640: in __call__
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m call(*args, call_id=call_id)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64))
        call       = <function ffi_call.<locals>.wrapped at 0x7f3be5946520>
        call_id    = 0
        d          = 0
        device     = 'cpu'
        i          = 1
        input_arg  = <warp._src.jax_experimental.ffi.FfiArg object at 0x7f3bdf1ce0c0>
        input_value = Array([[-0.16308578, -0.56740909,  0.93064292],
       [ 0.14900107,  0.06445298, -0.29018964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64)
        module     = <warp._src.context.Module object at 0x7f3c205b0fb0>
        num_inputs = 2
        out_types  = [ShapeDtypeStruct(shape=(4, 3), dtype=float64)]
        output_arg = <warp._src.jax_experimental.ffi.FfiArg object at 0x7f3bdf1cffb0>
        output_dims = (4, 3)
        self       = <warp._src.jax_experimental.ffi.FfiCallable object at 0x7f3bdf1cc620>
        static_inputs = {}
        vmap_method = 'broadcast_all'
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/ffi.py#x1B[0m:540: in wrapped
    #x1B[0mresults = ffi_call_p.bind(#x1B[90m#x1B[39;49;00m
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64))
        custom_call_api_version = 4
        has_side_effect = False
        in_avals   = [ShapedArray(float64[4,3]), ShapedArray(float64[4,3])]
        input_layouts = None
        input_output_aliases = {}
        kwargs     = {'call_id': 0}
        legacy_backend_config = None
        multiple_results = True
        output_layouts_ = None
        result_avals = (ShapedArray(float64[4,3]),)
        result_shape_dtypes = [ShapeDtypeStruct(shape=(4, 3), dtype=float64)]
        static_input_layouts = ((1, 0), (1, 0))
        static_input_output_aliases = ()
        static_output_layouts = ((1, 0),)
        target_name = 'WarpModelAdapter___hess_prod_callable__locals__hess_prod_callable_616'
        vmap_method = 'broadcast_all'
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:633: in bind
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m #x1B[96mself#x1B[39;49;00m._true_bind(*args, **params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = [Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64)]
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        self       = ffi_call
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:649: in _true_bind
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m #x1B[96mself#x1B[39;49;00m.bind_with_trace(prev_trace, args, params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        arg        = Array([[-0.16308578, -0.56740909,  0.93064292],
       [ 0.14900107,  0.06445298, -0.29018964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64)
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64))
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        prev_trace = EvalTrace
        self       = ffi_call
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:661: in bind_with_trace
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m trace.process_primitive(#x1B[96mself#x1B[39;49;00m, args, params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64))
        in_type    = [ShapedArray(float64[4,3]), ShapedArray(float64[4,3])]
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        self       = ffi_call
        trace      = EvalTrace
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:1210: in process_primitive
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m primitive.impl(*args, **params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = [Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64)]
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        primitive  = ffi_call
        self       = EvalTrace
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

prim = ffi_call
args = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64))
params = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
fun = <PjitFunction of <function ffi_call at 0x7f3be5944540>>
prev = <object object at 0x7f3ca25a3b00>

    #x1B[0m#x1B[94mdef#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[92mapply_primitive#x1B[39;49;00m(prim, *args, **params):#x1B[90m#x1B[39;49;00m
    #x1B[90m  #x1B[39;49;00m#x1B[33m"""Impl rule that compiles and runs a single primitive 'prim' using XLA."""#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
      fun = xla_primitive_callable(prim, **params)#x1B[90m#x1B[39;49;00m
      #x1B[90m# TODO(yashkatariya): Investigate adding is_primitive to jit and never#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
      #x1B[90m# triggering the disable jit path instead of messing around with it here.#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
      prev = config.disable_jit.swap_local(#x1B[94mFalse#x1B[39;49;00m)#x1B[90m#x1B[39;49;00m
      #x1B[94mtry#x1B[39;49;00m:#x1B[90m#x1B[39;49;00m
>       outs = fun(*args)#x1B[90m#x1B[39;49;00m
               ^^^^^^^^^^#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31mE       jax.errors.JaxRuntimeError: NOT_FOUND: No FFI handler registered for WarpModelAdapter___hess_prod_callable__locals__hess_prod_callable_616 on a platform Host (canonical host)#x1B[0m
#x1B[1m#x1B[31mE       --------------------#x1B[0m
#x1B[1m#x1B[31mE       For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.#x1B[0m
#x1B[1m#x1B[31mE       Falsifying example: test_arap_hess_prod(#x1B[0m
#x1B[1m#x1B[31mE           model=Model(dirichlet=Dirichlet(dim=3,#x1B[0m
#x1B[1m#x1B[31mE             dirichlet_index=Array([], shape=(0,), dtype=int64),#x1B[0m
#x1B[1m#x1B[31mE             dirichlet_value=Array([], shape=(0,), dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE             free_index=Array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11], dtype=int64),#x1B[0m
#x1B[1m#x1B[31mE             n_points=4),#x1B[0m
#x1B[1m#x1B[31mE            u_full=Array([[-0.23030316,  0.4469279 , -1.00805308],#x1B[0m
#x1B[1m#x1B[31mE                   [ 0.05685429,  0.38357654,  1.17551652],#x1B[0m
#x1B[1m#x1B[31mE                   [-0.32433812, -0.93279542,  0.88309023],#x1B[0m
#x1B[1m#x1B[31mE                   [-0.2684392 ,  0.67174205, -1.042468  ]], dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE            jax=JaxModel(energies={}),#x1B[0m
#x1B[1m#x1B[31mE            warp=WarpModelAdapter(wrapped=WarpModel(dim=3,#x1B[0m
#x1B[1m#x1B[31mE              energies={'elastic': Arap(id='elastic',#x1B[0m
#x1B[1m#x1B[31mE                requires_grad=['mu'],#x1B[0m
#x1B[1m#x1B[31mE                cells=array(shape=(1,), dtype=vec4i),#x1B[0m
#x1B[1m#x1B[31mE                dhdX=array(shape=(1, 1), dtype=mat43(d)),#x1B[0m
#x1B[1m#x1B[31mE                dV=array(shape=(1, 1), dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE                params=Arap__Params(#x1B[0m
#x1B[1m#x1B[31mE                	mu=array(shape=(1,), dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE                ),#x1B[0m
#x1B[1m#x1B[31mE                clamp_hess_diag=False,#x1B[0m
#x1B[1m#x1B[31mE                clamp_hess_quad=False,#x1B[0m
#x1B[1m#x1B[31mE                clamp_lambda=False)})),#x1B[0m
#x1B[1m#x1B[31mE            edges_length_mean=Array(0.99999998, dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE            frozen=False),#x1B[0m
#x1B[1m#x1B[31mE           mesh=UnstructuredGrid (0x7f3bf8d15f00)#x1B[0m
#x1B[1m#x1B[31mE             N Cells:    1#x1B[0m
#x1B[1m#x1B[31mE             N Points:   4#x1B[0m
#x1B[1m#x1B[31mE             X Bounds:   -3.536e-01, 3.536e-01#x1B[0m
#x1B[1m#x1B[31mE             Y Bounds:   -3.536e-01, 3.536e-01#x1B[0m
#x1B[1m#x1B[31mE             Z Bounds:   -3.536e-01, 3.536e-01#x1B[0m
#x1B[1m#x1B[31mE             N Arrays:   2,#x1B[0m
#x1B[1m#x1B[31mE           seed=0,  # or any other generated value#x1B[0m
#x1B[1m#x1B[31mE       )#x1B[0m
#x1B[1m#x1B[31mE       #x1B[0m
#x1B[1m#x1B[31mE       You can reproduce this example by temporarily adding @reproduce_failure('6.150.1', b'AEEA') as a decorator on your test case#x1B[0m

args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64))
fun        = <PjitFunction of <function ffi_call at 0x7f3be5944540>>
params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
prev       = <object object at 0x7f3ca25a3b00>
prim       = ffi_call

#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/dispatch.py#x1B[0m:91: JaxRuntimeError

tests/warp/energies/elastic/hyperelastic/test_arap.py::test_arap_hess_quad

Flake rate in main: 66.67% (Passed 19 times, Failed 38 times)

Stack Traces | 68.1s run time

model = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
mesh = UnstructuredGrid (0x7f3bf8d15f00)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   2

    #x1B[0m#x1B[37m@hypothesis#x1B[39;49;00m.given(seed=testing.seed())#x1B[90m#x1B[39;49;00m
>   #x1B[94mdef#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[92mtest_arap_hess_quad#x1B[39;49;00m(seed: #x1B[96mint#x1B[39;49;00m, model: Model, mesh: pv.UnstructuredGrid) -> #x1B[94mNone#x1B[39;49;00m:#x1B[90m#x1B[39;49;00m
                   ^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m

f          = <function given.<locals>.run_test_as_given.<locals>.wrapped_test at 0x7f3c401e87c0>
mesh       = UnstructuredGrid (0x7f3bf8d15f00)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   2
model      = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)

#x1B[1m#x1B[.../elastic/hyperelastic/test_arap.py#x1B[0m:64: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
#x1B[1m#x1B[.../elastic/hyperelastic/test_arap.py#x1B[0m:65: in test_arap_hess_quad
    #x1B[0mcommon.check_hess_quad(seed, model, mesh)#x1B[90m#x1B[39;49;00m
        mesh       = UnstructuredGrid (0x7f3bf8d15f00)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   2
        model      = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
        seed       = 0
#x1B[1m#x1B[.../elastic/hyperelastic/common.py#x1B[0m:62: in check_hess_quad
    #x1B[0mexpected: Scalar = jnp.vdot(p, model.hess_prod(u, p))#x1B[90m#x1B[39;49;00m
                                   ^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        actual     = Array(0.29022131, dtype=float64)
        mesh       = UnstructuredGrid (0x7f3bf8d15f00)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   2
        model      = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
        p          = Array([[-0.93533705, -0.11326748,  0.16791879],
       [ 0.81027622, -0.72888169,  0.62500202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)
        seed       = 0
        u          = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
#x1B[1m#x1B[.../apple/model/_model.py#x1B[0m:117: in hess_prod
    #x1B[0moutput_wp: Full = #x1B[96mself#x1B[39;49;00m.warp.hess_prod(u_full, p_full)#x1B[90m#x1B[39;49;00m
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        output_jax = Array([[0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.]], dtype=float64)
        p          = Array([[-0.93533705, -0.11326748,  0.16791879],
       [ 0.81027622, -0.72888169,  0.62500202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)
        p_full     = Array([[-0.93533705, -0.11326748,  0.16791879],
       [ 0.81027622, -0.72888169,  0.62500202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)
        self       = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
        u          = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
        u_full     = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
#x1B[1m#x1B[.../warp/model/_adapter.py#x1B[0m:60: in hess_prod
    #x1B[0m(output,) = #x1B[96mself#x1B[39;49;00m._hess_prod_callable(u, p, output_dims=u.shape)#x1B[90m#x1B[39;49;00m
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        p          = Array([[-0.93533705, -0.11326748,  0.16791879],
       [ 0.81027622, -0.72888169,  0.62500202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)
        self       = WarpModelAdapter(
  wrapped=WarpModel(
    energies={
      'elastic':
      Arap(
        id='elastic',
        requi...1,), dtype=float64),
               ),
        clamp_hess_diag=False,
        clamp_hess_quad=False
      )
    }
  )
)
        u          = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
#x1B[1m#x1B[31m.venv/lib/python3.12.../_src/jax_experimental/ffi.py#x1B[0m:640: in __call__
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m call(*args, call_id=call_id)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64))
        call       = <function ffi_call.<locals>.wrapped at 0x7f3be50bb2e0>
        call_id    = 0
        d          = 0
        device     = 'cpu'
        i          = 1
        input_arg  = <warp._src.jax_experimental.ffi.FfiArg object at 0x7f3be49292e0>
        input_value = Array([[-0.93533705, -0.11326748,  0.16791879],
       [ 0.81027622, -0.72888169,  0.62500202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)
        module     = <warp._src.context.Module object at 0x7f3c205b0fb0>
        num_inputs = 2
        out_types  = [ShapeDtypeStruct(shape=(4, 3), dtype=float64)]
        output_arg = <warp._src.jax_experimental.ffi.FfiArg object at 0x7f3be492bda0>
        output_dims = (4, 3)
        self       = <warp._src.jax_experimental.ffi.FfiCallable object at 0x7f3bdf42ab40>
        static_inputs = {}
        vmap_method = 'broadcast_all'
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/ffi.py#x1B[0m:540: in wrapped
    #x1B[0mresults = ffi_call_p.bind(#x1B[90m#x1B[39;49;00m
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64))
        custom_call_api_version = 4
        has_side_effect = False
        in_avals   = [ShapedArray(float64[4,3]), ShapedArray(float64[4,3])]
        input_layouts = None
        input_output_aliases = {}
        kwargs     = {'call_id': 0}
        legacy_backend_config = None
        multiple_results = True
        output_layouts_ = None
        result_avals = (ShapedArray(float64[4,3]),)
        result_shape_dtypes = [ShapeDtypeStruct(shape=(4, 3), dtype=float64)]
        static_input_layouts = ((1, 0), (1, 0))
        static_input_output_aliases = ()
        static_output_layouts = ((1, 0),)
        target_name = 'WarpModelAdapter___hess_prod_callable__locals__hess_prod_callable_515'
        vmap_method = 'broadcast_all'
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:633: in bind
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m #x1B[96mself#x1B[39;49;00m._true_bind(*args, **params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = [Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)]
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        self       = ffi_call
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:649: in _true_bind
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m #x1B[96mself#x1B[39;49;00m.bind_with_trace(prev_trace, args, params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        arg        = Array([[-0.93533705, -0.11326748,  0.16791879],
       [ 0.81027622, -0.72888169,  0.62500202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64))
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        prev_trace = EvalTrace
        self       = ffi_call
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:661: in bind_with_trace
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m trace.process_primitive(#x1B[96mself#x1B[39;49;00m, args, params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64))
        in_type    = [ShapedArray(float64[4,3]), ShapedArray(float64[4,3])]
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        self       = ffi_call
        trace      = EvalTrace
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:1210: in process_primitive
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m primitive.impl(*args, **params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = [Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)]
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        primitive  = ffi_call
        self       = EvalTrace
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

prim = ffi_call
args = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64))
params = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
fun = <PjitFunction of <function ffi_call at 0x7f3be50b9c60>>
prev = <object object at 0x7f3ca25a3b00>

    #x1B[0m#x1B[94mdef#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[92mapply_primitive#x1B[39;49;00m(prim, *args, **params):#x1B[90m#x1B[39;49;00m
    #x1B[90m  #x1B[39;49;00m#x1B[33m"""Impl rule that compiles and runs a single primitive 'prim' using XLA."""#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
      fun = xla_primitive_callable(prim, **params)#x1B[90m#x1B[39;49;00m
      #x1B[90m# TODO(yashkatariya): Investigate adding is_primitive to jit and never#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
      #x1B[90m# triggering the disable jit path instead of messing around with it here.#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
      prev = config.disable_jit.swap_local(#x1B[94mFalse#x1B[39;49;00m)#x1B[90m#x1B[39;49;00m
      #x1B[94mtry#x1B[39;49;00m:#x1B[90m#x1B[39;49;00m
>       outs = fun(*args)#x1B[90m#x1B[39;49;00m
               ^^^^^^^^^^#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31mE       jax.errors.JaxRuntimeError: NOT_FOUND: No FFI handler registered for WarpModelAdapter___hess_prod_callable__locals__hess_prod_callable_515 on a platform Host (canonical host)#x1B[0m
#x1B[1m#x1B[31mE       --------------------#x1B[0m
#x1B[1m#x1B[31mE       For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.#x1B[0m
#x1B[1m#x1B[31mE       Falsifying example: test_arap_hess_quad(#x1B[0m
#x1B[1m#x1B[31mE           model=Model(dirichlet=Dirichlet(dim=3,#x1B[0m
#x1B[1m#x1B[31mE             dirichlet_index=Array([], shape=(0,), dtype=int64),#x1B[0m
#x1B[1m#x1B[31mE             dirichlet_value=Array([], shape=(0,), dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE             free_index=Array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11], dtype=int64),#x1B[0m
#x1B[1m#x1B[31mE             n_points=4),#x1B[0m
#x1B[1m#x1B[31mE            u_full=Array([[ 0.47874223, -0.38700476,  0.24528764],#x1B[0m
#x1B[1m#x1B[31mE                   [ 0.38609176, -0.17329485,  0.90024778],#x1B[0m
#x1B[1m#x1B[31mE                   [ 0.81725247, -0.63050376, -0.88730332],#x1B[0m
#x1B[1m#x1B[31mE                   [-0.96864485, -0.64144405, -1.07241723]], dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE            jax=JaxModel(energies={}),#x1B[0m
#x1B[1m#x1B[31mE            warp=WarpModelAdapter(wrapped=WarpModel(dim=3,#x1B[0m
#x1B[1m#x1B[31mE              energies={'elastic': Arap(id='elastic',#x1B[0m
#x1B[1m#x1B[31mE                requires_grad=['mu'],#x1B[0m
#x1B[1m#x1B[31mE                cells=array(shape=(1,), dtype=vec4i),#x1B[0m
#x1B[1m#x1B[31mE                dhdX=array(shape=(1, 1), dtype=mat43(d)),#x1B[0m
#x1B[1m#x1B[31mE                dV=array(shape=(1, 1), dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE                params=Arap__Params(#x1B[0m
#x1B[1m#x1B[31mE                	mu=array(shape=(1,), dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE                ),#x1B[0m
#x1B[1m#x1B[31mE                clamp_hess_diag=False,#x1B[0m
#x1B[1m#x1B[31mE                clamp_hess_quad=False,#x1B[0m
#x1B[1m#x1B[31mE                clamp_lambda=False)})),#x1B[0m
#x1B[1m#x1B[31mE            edges_length_mean=Array(0.99999998, dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE            frozen=False),#x1B[0m
#x1B[1m#x1B[31mE           mesh=UnstructuredGrid (0x7f3bf8d15f00)#x1B[0m
#x1B[1m#x1B[31mE             N Cells:    1#x1B[0m
#x1B[1m#x1B[31mE             N Points:   4#x1B[0m
#x1B[1m#x1B[31mE             X Bounds:   -3.536e-01, 3.536e-01#x1B[0m
#x1B[1m#x1B[31mE             Y Bounds:   -3.536e-01, 3.536e-01#x1B[0m
#x1B[1m#x1B[31mE             Z Bounds:   -3.536e-01, 3.536e-01#x1B[0m
#x1B[1m#x1B[31mE             N Arrays:   2,#x1B[0m
#x1B[1m#x1B[31mE           seed=0,  # or any other generated value#x1B[0m
#x1B[1m#x1B[31mE       )#x1B[0m
#x1B[1m#x1B[31mE       #x1B[0m
#x1B[1m#x1B[31mE       You can reproduce this example by temporarily adding @reproduce_failure('6.150.1', b'AEEA') as a decorator on your test case#x1B[0m

args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64))
fun        = <PjitFunction of <function ffi_call at 0x7f3be50b9c60>>
params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
prev       = <object object at 0x7f3ca25a3b00>
prim       = ffi_call

#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/dispatch.py#x1B[0m:91: JaxRuntimeError

tests/warp/energies/elastic/hyperelastic/test_arap_muscle.py::test_arap_muscle_hess_prod

Flake rate in main: 71.15% (Passed 15 times, Failed 37 times)

Stack Traces | 60.2s run time

model = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
mesh = UnstructuredGrid (0x7f3bedc1dea0)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   3

    #x1B[0m#x1B[37m@hypothesis#x1B[39;49;00m.given(seed=testing.seed())#x1B[90m#x1B[39;49;00m
>   #x1B[94mdef#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[92mtest_arap_muscle_hess_prod#x1B[39;49;00m(#x1B[90m#x1B[39;49;00m
                   ^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        seed: #x1B[96mint#x1B[39;49;00m, model: Model, mesh: pv.UnstructuredGrid#x1B[90m#x1B[39;49;00m
    ) -> #x1B[94mNone#x1B[39;49;00m:#x1B[90m#x1B[39;49;00m

f          = <function given.<locals>.run_test_as_given.<locals>.wrapped_test at 0x7f3c401ea3e0>
mesh       = UnstructuredGrid (0x7f3bedc1dea0)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   3
model      = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)

#x1B[1m#x1B[.../elastic/hyperelastic/test_arap_muscle.py#x1B[0m:64: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
#x1B[1m#x1B[.../elastic/hyperelastic/test_arap_muscle.py#x1B[0m:67: in test_arap_muscle_hess_prod
    #x1B[0mcommon.check_hess_prod(seed, model, mesh)#x1B[90m#x1B[39;49;00m
        mesh       = UnstructuredGrid (0x7f3bedc1dea0)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   3
        model      = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
        seed       = 0
#x1B[1m#x1B[.../elastic/hyperelastic/common.py#x1B[0m:55: in check_hess_prod
    #x1B[0mtesting.check_jvp(model.grad, model.hess_prod, u, rtol=#x1B[94m1e-3#x1B[39;49;00m)#x1B[90m#x1B[39;49;00m
        mesh       = UnstructuredGrid (0x7f3bedc1dea0)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   3
        model      = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
        seed       = 0
        u          = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
#x1B[1m#x1B[.../apple/model/_model.py#x1B[0m:117: in hess_prod
    #x1B[0moutput_wp: Full = #x1B[96mself#x1B[39;49;00m.warp.hess_prod(u_full, p_full)#x1B[90m#x1B[39;49;00m
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        output_jax = Array([[0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.]], dtype=float64)
        p          = Array([[-0.16308578, -0.56740909,  0.93064292],
       [ 0.14900107,  0.06445298, -0.29018964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64)
        p_full     = Array([[-0.16308578, -0.56740909,  0.93064292],
       [ 0.14900107,  0.06445298, -0.29018964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64)
        self       = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
        u          = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
        u_full     = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
#x1B[1m#x1B[.../warp/model/_adapter.py#x1B[0m:60: in hess_prod
    #x1B[0m(output,) = #x1B[96mself#x1B[39;49;00m._hess_prod_callable(u, p, output_dims=u.shape)#x1B[90m#x1B[39;49;00m
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        p          = Array([[-0.16308578, -0.56740909,  0.93064292],
       [ 0.14900107,  0.06445298, -0.29018964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64)
        self       = WarpModelAdapter(
  wrapped=WarpModel(
    energies={
      'elastic':
      ArapMuscle(
        id='elastic',
       ...1,), dtype=float64),
               ),
        clamp_hess_diag=False,
        clamp_hess_quad=False
      )
    }
  )
)
        u          = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
#x1B[1m#x1B[31m.venv/lib/python3.12.../_src/jax_experimental/ffi.py#x1B[0m:640: in __call__
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m call(*args, call_id=call_id)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64))
        call       = <function ffi_call.<locals>.wrapped at 0x7f3bec239580>
        call_id    = 0
        d          = 0
        device     = 'cpu'
        i          = 1
        input_arg  = <warp._src.jax_experimental.ffi.FfiArg object at 0x7f3be5073bf0>
        input_value = Array([[-0.16308578, -0.56740909,  0.93064292],
       [ 0.14900107,  0.06445298, -0.29018964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64)
        module     = <warp._src.context.Module object at 0x7f3c205b0fb0>
        num_inputs = 2
        out_types  = [ShapeDtypeStruct(shape=(4, 3), dtype=float64)]
        output_arg = <warp._src.jax_experimental.ffi.FfiArg object at 0x7f3be5070fe0>
        output_dims = (4, 3)
        self       = <warp._src.jax_experimental.ffi.FfiCallable object at 0x7f3be5072180>
        static_inputs = {}
        vmap_method = 'broadcast_all'
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/ffi.py#x1B[0m:540: in wrapped
    #x1B[0mresults = ffi_call_p.bind(#x1B[90m#x1B[39;49;00m
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64))
        custom_call_api_version = 4
        has_side_effect = False
        in_avals   = [ShapedArray(float64[4,3]), ShapedArray(float64[4,3])]
        input_layouts = None
        input_output_aliases = {}
        kwargs     = {'call_id': 0}
        legacy_backend_config = None
        multiple_results = True
        output_layouts_ = None
        result_avals = (ShapedArray(float64[4,3]),)
        result_shape_dtypes = [ShapeDtypeStruct(shape=(4, 3), dtype=float64)]
        static_input_layouts = ((1, 0), (1, 0))
        static_input_output_aliases = ()
        static_output_layouts = ((1, 0),)
        target_name = 'WarpModelAdapter___hess_prod_callable__locals__hess_prod_callable_718'
        vmap_method = 'broadcast_all'
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:633: in bind
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m #x1B[96mself#x1B[39;49;00m._true_bind(*args, **params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = [Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64)]
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        self       = ffi_call
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:649: in _true_bind
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m #x1B[96mself#x1B[39;49;00m.bind_with_trace(prev_trace, args, params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        arg        = Array([[-0.16308578, -0.56740909,  0.93064292],
       [ 0.14900107,  0.06445298, -0.29018964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64)
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64))
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        prev_trace = EvalTrace
        self       = ffi_call
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:661: in bind_with_trace
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m trace.process_primitive(#x1B[96mself#x1B[39;49;00m, args, params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64))
        in_type    = [ShapedArray(float64[4,3]), ShapedArray(float64[4,3])]
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        self       = ffi_call
        trace      = EvalTrace
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:1210: in process_primitive
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m primitive.impl(*args, **params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = [Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64)]
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        primitive  = ffi_call
        self       = EvalTrace
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

prim = ffi_call
args = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64))
params = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
fun = <PjitFunction of <function ffi_call at 0x7f3bec2398a0>>
prev = <object object at 0x7f3ca25a3b00>

    #x1B[0m#x1B[94mdef#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[92mapply_primitive#x1B[39;49;00m(prim, *args, **params):#x1B[90m#x1B[39;49;00m
    #x1B[90m  #x1B[39;49;00m#x1B[33m"""Impl rule that compiles and runs a single primitive 'prim' using XLA."""#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
      fun = xla_primitive_callable(prim, **params)#x1B[90m#x1B[39;49;00m
      #x1B[90m# TODO(yashkatariya): Investigate adding is_primitive to jit and never#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
      #x1B[90m# triggering the disable jit path instead of messing around with it here.#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
      prev = config.disable_jit.swap_local(#x1B[94mFalse#x1B[39;49;00m)#x1B[90m#x1B[39;49;00m
      #x1B[94mtry#x1B[39;49;00m:#x1B[90m#x1B[39;49;00m
>       outs = fun(*args)#x1B[90m#x1B[39;49;00m
               ^^^^^^^^^^#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31mE       jax.errors.JaxRuntimeError: NOT_FOUND: No FFI handler registered for WarpModelAdapter___hess_prod_callable__locals__hess_prod_callable_718 on a platform Host (canonical host)#x1B[0m
#x1B[1m#x1B[31mE       --------------------#x1B[0m
#x1B[1m#x1B[31mE       For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.#x1B[0m
#x1B[1m#x1B[31mE       Falsifying example: test_arap_muscle_hess_prod(#x1B[0m
#x1B[1m#x1B[31mE           model=Model(dirichlet=Dirichlet(dim=3,#x1B[0m
#x1B[1m#x1B[31mE             dirichlet_index=Array([], shape=(0,), dtype=int64),#x1B[0m
#x1B[1m#x1B[31mE             dirichlet_value=Array([], shape=(0,), dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE             free_index=Array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11], dtype=int64),#x1B[0m
#x1B[1m#x1B[31mE             n_points=4),#x1B[0m
#x1B[1m#x1B[31mE            u_full=Array([[-1.19693091,  0.13856566,  1.22283071],#x1B[0m
#x1B[1m#x1B[31mE                   [ 0.5109477 ,  0.2886049 ,  0.25589075],#x1B[0m
#x1B[1m#x1B[31mE                   [ 0.31821538, -0.3961438 , -0.7275457 ],#x1B[0m
#x1B[1m#x1B[31mE                   [-0.93225488, -0.42896687,  0.04529193]], dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE            jax=JaxModel(energies={}),#x1B[0m
#x1B[1m#x1B[31mE            warp=WarpModelAdapter(wrapped=WarpModel(dim=3,#x1B[0m
#x1B[1m#x1B[31mE              energies={'elastic': ArapMuscle(id='elastic',#x1B[0m
#x1B[1m#x1B[31mE                requires_grad=['activation', 'mu'],#x1B[0m
#x1B[1m#x1B[31mE                cells=array(shape=(1,), dtype=vec4i),#x1B[0m
#x1B[1m#x1B[31mE                dhdX=array(shape=(1, 1), dtype=mat43(d)),#x1B[0m
#x1B[1m#x1B[31mE                dV=array(shape=(1, 1), dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE                params=ArapMuscle__Params(#x1B[0m
#x1B[1m#x1B[31mE                	activation=array(shape=(1,), dtype=vector(length=6, dtype=float64)),#x1B[0m
#x1B[1m#x1B[31mE                	mu=array(shape=(1,), dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE                ),#x1B[0m
#x1B[1m#x1B[31mE                clamp_hess_diag=False,#x1B[0m
#x1B[1m#x1B[31mE                clamp_hess_quad=False,#x1B[0m
#x1B[1m#x1B[31mE                clamp_lambda=False)})),#x1B[0m
#x1B[1m#x1B[31mE            edges_length_mean=Array(0.99999998, dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE            frozen=False),#x1B[0m
#x1B[1m#x1B[31mE           mesh=UnstructuredGrid (0x7f3bedc1dea0)#x1B[0m
#x1B[1m#x1B[31mE             N Cells:    1#x1B[0m
#x1B[1m#x1B[31mE             N Points:   4#x1B[0m
#x1B[1m#x1B[31mE             X Bounds:   -3.536e-01, 3.536e-01#x1B[0m
#x1B[1m#x1B[31mE             Y Bounds:   -3.536e-01, 3.536e-01#x1B[0m
#x1B[1m#x1B[31mE             Z Bounds:   -3.536e-01, 3.536e-01#x1B[0m
#x1B[1m#x1B[31mE             N Arrays:   3,#x1B[0m
#x1B[1m#x1B[31mE           seed=0,  # or any other generated value#x1B[0m
#x1B[1m#x1B[31mE       )#x1B[0m
#x1B[1m#x1B[31mE       #x1B[0m
#x1B[1m#x1B[31mE       You can reproduce this example by temporarily adding @reproduce_failure('6.150.1', b'AEEA') as a decorator on your test case#x1B[0m

args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64))
fun        = <PjitFunction of <function ffi_call at 0x7f3bec2398a0>>
params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
prev       = <object object at 0x7f3ca25a3b00>
prim       = ffi_call

#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/dispatch.py#x1B[0m:91: JaxRuntimeError

tests/warp/energies/elastic/hyperelastic/test_arap_muscle.py::test_arap_muscle_hess_quad

Flake rate in main: 71.15% (Passed 15 times, Failed 37 times)

Stack Traces | 66.4s run time

model = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
mesh = UnstructuredGrid (0x7f3bedc1dea0)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   3

    #x1B[0m#x1B[37m@hypothesis#x1B[39;49;00m.given(seed=testing.seed())#x1B[90m#x1B[39;49;00m
>   #x1B[94mdef#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[92mtest_arap_muscle_hess_quad#x1B[39;49;00m(#x1B[90m#x1B[39;49;00m
                   ^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        seed: #x1B[96mint#x1B[39;49;00m, model: Model, mesh: pv.UnstructuredGrid#x1B[90m#x1B[39;49;00m
    ) -> #x1B[94mNone#x1B[39;49;00m:#x1B[90m#x1B[39;49;00m

f          = <function given.<locals>.run_test_as_given.<locals>.wrapped_test at 0x7f3c401ea840>
mesh       = UnstructuredGrid (0x7f3bedc1dea0)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   3
model      = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)

#x1B[1m#x1B[.../elastic/hyperelastic/test_arap_muscle.py#x1B[0m:71: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
#x1B[1m#x1B[.../elastic/hyperelastic/test_arap_muscle.py#x1B[0m:74: in test_arap_muscle_hess_quad
    #x1B[0mcommon.check_hess_quad(seed, model, mesh)#x1B[90m#x1B[39;49;00m
        mesh       = UnstructuredGrid (0x7f3bedc1dea0)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   3
        model      = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
        seed       = 0
#x1B[1m#x1B[.../elastic/hyperelastic/common.py#x1B[0m:62: in check_hess_quad
    #x1B[0mexpected: Scalar = jnp.vdot(p, model.hess_prod(u, p))#x1B[90m#x1B[39;49;00m
                                   ^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        actual     = Array(0.35981608, dtype=float64)
        mesh       = UnstructuredGrid (0x7f3bedc1dea0)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   3
        model      = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
        p          = Array([[-0.93533705, -0.11326748,  0.16791879],
       [ 0.81027622, -0.72888169,  0.62500202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)
        seed       = 0
        u          = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
#x1B[1m#x1B[.../apple/model/_model.py#x1B[0m:117: in hess_prod
    #x1B[0moutput_wp: Full = #x1B[96mself#x1B[39;49;00m.warp.hess_prod(u_full, p_full)#x1B[90m#x1B[39;49;00m
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        output_jax = Array([[0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.]], dtype=float64)
        p          = Array([[-0.93533705, -0.11326748,  0.16791879],
       [ 0.81027622, -0.72888169,  0.62500202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)
        p_full     = Array([[-0.93533705, -0.11326748,  0.16791879],
       [ 0.81027622, -0.72888169,  0.62500202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)
        self       = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
        u          = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
        u_full     = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
#x1B[1m#x1B[.../warp/model/_adapter.py#x1B[0m:60: in hess_prod
    #x1B[0m(output,) = #x1B[96mself#x1B[39;49;00m._hess_prod_callable(u, p, output_dims=u.shape)#x1B[90m#x1B[39;49;00m
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        p          = Array([[-0.93533705, -0.11326748,  0.16791879],
       [ 0.81027622, -0.72888169,  0.62500202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)
        self       = WarpModelAdapter(
  wrapped=WarpModel(
    energies={
      'elastic':
      ArapMuscle(
        id='elastic',
       ...1,), dtype=float64),
               ),
        clamp_hess_diag=False,
        clamp_hess_quad=False
      )
    }
  )
)
        u          = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
#x1B[1m#x1B[31m.venv/lib/python3.12.../_src/jax_experimental/ffi.py#x1B[0m:640: in __call__
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m call(*args, call_id=call_id)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64))
        call       = <function ffi_call.<locals>.wrapped at 0x7f3be50b8cc0>
        call_id    = 0
        d          = 0
        device     = 'cpu'
        i          = 1
        input_arg  = <warp._src.jax_experimental.ffi.FfiArg object at 0x7f3bec9ce960>
        input_value = Array([[-0.93533705, -0.11326748,  0.16791879],
       [ 0.81027622, -0.72888169,  0.62500202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)
        module     = <warp._src.context.Module object at 0x7f3c205b0fb0>
        num_inputs = 2
        out_types  = [ShapeDtypeStruct(shape=(4, 3), dtype=float64)]
        output_arg = <warp._src.jax_experimental.ffi.FfiArg object at 0x7f3bec996540>
        output_dims = (4, 3)
        self       = <warp._src.jax_experimental.ffi.FfiCallable object at 0x7f3be5278b00>
        static_inputs = {}
        vmap_method = 'broadcast_all'
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/ffi.py#x1B[0m:540: in wrapped
    #x1B[0mresults = ffi_call_p.bind(#x1B[90m#x1B[39;49;00m
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64))
        custom_call_api_version = 4
        has_side_effect = False
        in_avals   = [ShapedArray(float64[4,3]), ShapedArray(float64[4,3])]
        input_layouts = None
        input_output_aliases = {}
        kwargs     = {'call_id': 0}
        legacy_backend_config = None
        multiple_results = True
        output_layouts_ = None
        result_avals = (ShapedArray(float64[4,3]),)
        result_shape_dtypes = [ShapeDtypeStruct(shape=(4, 3), dtype=float64)]
        static_input_layouts = ((1, 0), (1, 0))
        static_input_output_aliases = ()
        static_output_layouts = ((1, 0),)
        target_name = 'WarpModelAdapter___hess_prod_callable__locals__hess_prod_callable_824'
        vmap_method = 'broadcast_all'
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:633: in bind
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m #x1B[96mself#x1B[39;49;00m._true_bind(*args, **params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = [Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)]
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        self       = ffi_call
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:649: in _true_bind
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m #x1B[96mself#x1B[39;49;00m.bind_with_trace(prev_trace, args, params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        arg        = Array([[-0.93533705, -0.11326748,  0.16791879],
       [ 0.81027622, -0.72888169,  0.62500202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64))
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        prev_trace = EvalTrace
        self       = ffi_call
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:661: in bind_with_trace
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m trace.process_primitive(#x1B[96mself#x1B[39;49;00m, args, params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64))
        in_type    = [ShapedArray(float64[4,3]), ShapedArray(float64[4,3])]
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        self       = ffi_call
        trace      = EvalTrace
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:1210: in process_primitive
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m primitive.impl(*args, **params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = [Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)]
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        primitive  = ffi_call
        self       = EvalTrace
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

prim = ffi_call
args = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64))
params = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
fun = <PjitFunction of <function ffi_call at 0x7f3be50bbf60>>
prev = <object object at 0x7f3ca25a3b00>

    #x1B[0m#x1B[94mdef#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[92mapply_primitive#x1B[39;49;00m(prim, *args, **params):#x1B[90m#x1B[39;49;00m
    #x1B[90m  #x1B[39;49;00m#x1B[33m"""Impl rule that compiles and runs a single primitive 'prim' using XLA."""#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
      fun = xla_primitive_callable(prim, **params)#x1B[90m#x1B[39;49;00m
      #x1B[90m# TODO(yashkatariya): Investigate adding is_primitive to jit and never#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
      #x1B[90m# triggering the disable jit path instead of messing around with it here.#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
      prev = config.disable_jit.swap_local(#x1B[94mFalse#x1B[39;49;00m)#x1B[90m#x1B[39;49;00m
      #x1B[94mtry#x1B[39;49;00m:#x1B[90m#x1B[39;49;00m
>       outs = fun(*args)#x1B[90m#x1B[39;49;00m
               ^^^^^^^^^^#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31mE       jax.errors.JaxRuntimeError: NOT_FOUND: No FFI handler registered for WarpModelAdapter___hess_prod_callable__locals__hess_prod_callable_824 on a platform Host (canonical host)#x1B[0m
#x1B[1m#x1B[31mE       --------------------#x1B[0m
#x1B[1m#x1B[31mE       For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.#x1B[0m
#x1B[1m#x1B[31mE       Falsifying example: test_arap_muscle_hess_quad(#x1B[0m
#x1B[1m#x1B[31mE           model=Model(dirichlet=Dirichlet(dim=3,#x1B[0m
#x1B[1m#x1B[31mE             dirichlet_index=Array([], shape=(0,), dtype=int64),#x1B[0m
#x1B[1m#x1B[31mE             dirichlet_value=Array([], shape=(0,), dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE             free_index=Array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11], dtype=int64),#x1B[0m
#x1B[1m#x1B[31mE             n_points=4),#x1B[0m
#x1B[1m#x1B[31mE            u_full=Array([[-0.76655026, -0.62556086,  0.68767497],#x1B[0m
#x1B[1m#x1B[31mE                   [-0.28233729, -0.69772602, -0.37828513],#x1B[0m
#x1B[1m#x1B[31mE                   [-0.04823423, -1.03855213,  0.95418263],#x1B[0m
#x1B[1m#x1B[31mE                   [ 1.07699311,  0.54753535, -0.98902473]], dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE            jax=JaxModel(energies={}),#x1B[0m
#x1B[1m#x1B[31mE            warp=WarpModelAdapter(wrapped=WarpModel(dim=3,#x1B[0m
#x1B[1m#x1B[31mE              energies={'elastic': ArapMuscle(id='elastic',#x1B[0m
#x1B[1m#x1B[31mE                requires_grad=['activation', 'mu'],#x1B[0m
#x1B[1m#x1B[31mE                cells=array(shape=(1,), dtype=vec4i),#x1B[0m
#x1B[1m#x1B[31mE                dhdX=array(shape=(1, 1), dtype=mat43(d)),#x1B[0m
#x1B[1m#x1B[31mE                dV=array(shape=(1, 1), dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE                params=ArapMuscle__Params(#x1B[0m
#x1B[1m#x1B[31mE                	activation=array(shape=(1,), dtype=vector(length=6, dtype=float64)),#x1B[0m
#x1B[1m#x1B[31mE                	mu=array(shape=(1,), dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE                ),#x1B[0m
#x1B[1m#x1B[31mE                clamp_hess_diag=False,#x1B[0m
#x1B[1m#x1B[31mE                clamp_hess_quad=False,#x1B[0m
#x1B[1m#x1B[31mE                clamp_lambda=False)})),#x1B[0m
#x1B[1m#x1B[31mE            edges_length_mean=Array(0.99999998, dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE            frozen=False),#x1B[0m
#x1B[1m#x1B[31mE           mesh=UnstructuredGrid (0x7f3bedc1dea0)#x1B[0m
#x1B[1m#x1B[31mE             N Cells:    1#x1B[0m
#x1B[1m#x1B[31mE             N Points:   4#x1B[0m
#x1B[1m#x1B[31mE             X Bounds:   -3.536e-01, 3.536e-01#x1B[0m
#x1B[1m#x1B[31mE             Y Bounds:   -3.536e-01, 3.536e-01#x1B[0m
#x1B[1m#x1B[31mE             Z Bounds:   -3.536e-01, 3.536e-01#x1B[0m
#x1B[1m#x1B[31mE             N Arrays:   3,#x1B[0m
#x1B[1m#x1B[31mE           seed=0,  # or any other generated value#x1B[0m
#x1B[1m#x1B[31mE       )#x1B[0m
#x1B[1m#x1B[31mE       #x1B[0m
#x1B[1m#x1B[31mE       You can reproduce this example by temporarily adding @reproduce_failure('6.150.1', b'AEEA') as a decorator on your test case#x1B[0m

args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64))
fun        = <PjitFunction of <function ffi_call at 0x7f3be50bbf60>>
params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
prev       = <object object at 0x7f3ca25a3b00>
prim       = ffi_call

#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/dispatch.py#x1B[0m:91: JaxRuntimeError

tests/warp/energies/elastic/hyperelastic/test_arap_muscle_v2.py::test_arap_muscle_hess_prod

Flake rate in main: 100.00% (Passed 0 times, Failed 15 times)

Stack Traces | 63.5s run time

model = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
mesh = UnstructuredGrid (0x7f3bdf179cc0)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   3

    #x1B[0m#x1B[37m@hypothesis#x1B[39;49;00m.given(seed=testing.seed())#x1B[90m#x1B[39;49;00m
>   #x1B[94mdef#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[92mtest_arap_muscle_hess_prod#x1B[39;49;00m(#x1B[90m#x1B[39;49;00m
                   ^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        seed: #x1B[96mint#x1B[39;49;00m, model: Model, mesh: pv.UnstructuredGrid#x1B[90m#x1B[39;49;00m
    ) -> #x1B[94mNone#x1B[39;49;00m:#x1B[90m#x1B[39;49;00m

f          = <function given.<locals>.run_test_as_given.<locals>.wrapped_test at 0x7f3c401b4a40>
mesh       = UnstructuredGrid (0x7f3bdf179cc0)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   3
model      = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)

#x1B[1m#x1B[.../elastic/hyperelastic/test_arap_muscle_v2.py#x1B[0m:64: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
#x1B[1m#x1B[.../elastic/hyperelastic/test_arap_muscle_v2.py#x1B[0m:67: in test_arap_muscle_hess_prod
    #x1B[0mcommon.check_hess_prod(seed, model, mesh)#x1B[90m#x1B[39;49;00m
        mesh       = UnstructuredGrid (0x7f3bdf179cc0)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   3
        model      = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
        seed       = 0
#x1B[1m#x1B[.../elastic/hyperelastic/common.py#x1B[0m:55: in check_hess_prod
    #x1B[0mtesting.check_jvp(model.grad, model.hess_prod, u, rtol=#x1B[94m1e-3#x1B[39;49;00m)#x1B[90m#x1B[39;49;00m
        mesh       = UnstructuredGrid (0x7f3bdf179cc0)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   3
        model      = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
        seed       = 0
        u          = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
#x1B[1m#x1B[.../apple/model/_model.py#x1B[0m:117: in hess_prod
    #x1B[0moutput_wp: Full = #x1B[96mself#x1B[39;49;00m.warp.hess_prod(u_full, p_full)#x1B[90m#x1B[39;49;00m
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        output_jax = Array([[0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.]], dtype=float64)
        p          = Array([[-0.16308578, -0.56740909,  0.93064292],
       [ 0.14900107,  0.06445298, -0.29018964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64)
        p_full     = Array([[-0.16308578, -0.56740909,  0.93064292],
       [ 0.14900107,  0.06445298, -0.29018964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64)
        self       = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
        u          = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
        u_full     = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
#x1B[1m#x1B[.../warp/model/_adapter.py#x1B[0m:60: in hess_prod
    #x1B[0m(output,) = #x1B[96mself#x1B[39;49;00m._hess_prod_callable(u, p, output_dims=u.shape)#x1B[90m#x1B[39;49;00m
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        p          = Array([[-0.16308578, -0.56740909,  0.93064292],
       [ 0.14900107,  0.06445298, -0.29018964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64)
        self       = WarpModelAdapter(
  wrapped=WarpModel(
    energies={
      'elastic':
      ArapMuscleV2(
        id='elastic',
     ...1,), dtype=float64),
               ),
        clamp_hess_diag=False,
        clamp_hess_quad=False
      )
    }
  )
)
        u          = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
#x1B[1m#x1B[31m.venv/lib/python3.12.../_src/jax_experimental/ffi.py#x1B[0m:640: in __call__
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m call(*args, call_id=call_id)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64))
        call       = <function ffi_call.<locals>.wrapped at 0x7f3be73211c0>
        call_id    = 0
        d          = 0
        device     = 'cpu'
        i          = 1
        input_arg  = <warp._src.jax_experimental.ffi.FfiArg object at 0x7f3bdf5eb9b0>
        input_value = Array([[-0.16308578, -0.56740909,  0.93064292],
       [ 0.14900107,  0.06445298, -0.29018964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64)
        module     = <warp._src.context.Module object at 0x7f3c205b0fb0>
        num_inputs = 2
        out_types  = [ShapeDtypeStruct(shape=(4, 3), dtype=float64)]
        output_arg = <warp._src.jax_experimental.ffi.FfiArg object at 0x7f3bdf5e9550>
        output_dims = (4, 3)
        self       = <warp._src.jax_experimental.ffi.FfiCallable object at 0x7f3bdf080530>
        static_inputs = {}
        vmap_method = 'broadcast_all'
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/ffi.py#x1B[0m:540: in wrapped
    #x1B[0mresults = ffi_call_p.bind(#x1B[90m#x1B[39;49;00m
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64))
        custom_call_api_version = 4
        has_side_effect = False
        in_avals   = [ShapedArray(float64[4,3]), ShapedArray(float64[4,3])]
        input_layouts = None
        input_output_aliases = {}
        kwargs     = {'call_id': 0}
        legacy_backend_config = None
        multiple_results = True
        output_layouts_ = None
        result_avals = (ShapedArray(float64[4,3]),)
        result_shape_dtypes = [ShapeDtypeStruct(shape=(4, 3), dtype=float64)]
        static_input_layouts = ((1, 0), (1, 0))
        static_input_output_aliases = ()
        static_output_layouts = ((1, 0),)
        target_name = 'WarpModelAdapter___hess_prod_callable__locals__hess_prod_callable_418'
        vmap_method = 'broadcast_all'
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:633: in bind
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m #x1B[96mself#x1B[39;49;00m._true_bind(*args, **params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = [Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64)]
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        self       = ffi_call
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:649: in _true_bind
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m #x1B[96mself#x1B[39;49;00m.bind_with_trace(prev_trace, args, params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        arg        = Array([[-0.16308578, -0.56740909,  0.93064292],
       [ 0.14900107,  0.06445298, -0.29018964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64)
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64))
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        prev_trace = EvalTrace
        self       = ffi_call
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:661: in bind_with_trace
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m trace.process_primitive(#x1B[96mself#x1B[39;49;00m, args, params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64))
        in_type    = [ShapedArray(float64[4,3]), ShapedArray(float64[4,3])]
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        self       = ffi_call
        trace      = EvalTrace
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:1210: in process_primitive
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m primitive.impl(*args, **params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = [Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64)]
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        primitive  = ffi_call
        self       = EvalTrace
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

prim = ffi_call
args = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64))
params = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
fun = <PjitFunction of <function ffi_call at 0x7f3be7320a40>>
prev = <object object at 0x7f3ca25a3b00>

    #x1B[0m#x1B[94mdef#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[92mapply_primitive#x1B[39;49;00m(prim, *args, **params):#x1B[90m#x1B[39;49;00m
    #x1B[90m  #x1B[39;49;00m#x1B[33m"""Impl rule that compiles and runs a single primitive 'prim' using XLA."""#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
      fun = xla_primitive_callable(prim, **params)#x1B[90m#x1B[39;49;00m
      #x1B[90m# TODO(yashkatariya): Investigate adding is_primitive to jit and never#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
      #x1B[90m# triggering the disable jit path instead of messing around with it here.#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
      prev = config.disable_jit.swap_local(#x1B[94mFalse#x1B[39;49;00m)#x1B[90m#x1B[39;49;00m
      #x1B[94mtry#x1B[39;49;00m:#x1B[90m#x1B[39;49;00m
>       outs = fun(*args)#x1B[90m#x1B[39;49;00m
               ^^^^^^^^^^#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31mE       jax.errors.JaxRuntimeError: NOT_FOUND: No FFI handler registered for WarpModelAdapter___hess_prod_callable__locals__hess_prod_callable_418 on a platform Host (canonical host)#x1B[0m
#x1B[1m#x1B[31mE       --------------------#x1B[0m
#x1B[1m#x1B[31mE       For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.#x1B[0m
#x1B[1m#x1B[31mE       Falsifying example: test_arap_muscle_hess_prod(#x1B[0m
#x1B[1m#x1B[31mE           model=Model(dirichlet=Dirichlet(dim=3,#x1B[0m
#x1B[1m#x1B[31mE             dirichlet_index=Array([], shape=(0,), dtype=int64),#x1B[0m
#x1B[1m#x1B[31mE             dirichlet_value=Array([], shape=(0,), dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE             free_index=Array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11], dtype=int64),#x1B[0m
#x1B[1m#x1B[31mE             n_points=4),#x1B[0m
#x1B[1m#x1B[31mE            u_full=Array([[-1.19693091,  0.13856566,  1.22283071],#x1B[0m
#x1B[1m#x1B[31mE                   [ 0.5109477 ,  0.2886049 ,  0.25589075],#x1B[0m
#x1B[1m#x1B[31mE                   [ 0.31821538, -0.3961438 , -0.7275457 ],#x1B[0m
#x1B[1m#x1B[31mE                   [-0.93225488, -0.42896687,  0.04529193]], dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE            jax=JaxModel(energies={}),#x1B[0m
#x1B[1m#x1B[31mE            warp=WarpModelAdapter(wrapped=WarpModel(dim=3,#x1B[0m
#x1B[1m#x1B[31mE              energies={'elastic': ArapMuscleV2(id='elastic',#x1B[0m
#x1B[1m#x1B[31mE                requires_grad=['activation', 'mu'],#x1B[0m
#x1B[1m#x1B[31mE                cells=array(shape=(1,), dtype=vec4i),#x1B[0m
#x1B[1m#x1B[31mE                dhdX=array(shape=(1, 1), dtype=mat43(d)),#x1B[0m
#x1B[1m#x1B[31mE                dV=array(shape=(1, 1), dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE                params=ArapMuscleV2__Params(#x1B[0m
#x1B[1m#x1B[31mE                	activation=array(shape=(1,), dtype=vector(length=6, dtype=float64)),#x1B[0m
#x1B[1m#x1B[31mE                	mu=array(shape=(1,), dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE                ),#x1B[0m
#x1B[1m#x1B[31mE                clamp_hess_diag=False,#x1B[0m
#x1B[1m#x1B[31mE                clamp_hess_quad=False,#x1B[0m
#x1B[1m#x1B[31mE                clamp_lambda=False)})),#x1B[0m
#x1B[1m#x1B[31mE            edges_length_mean=Array(0.99999998, dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE            frozen=False),#x1B[0m
#x1B[1m#x1B[31mE           mesh=UnstructuredGrid (0x7f3bdf179cc0)#x1B[0m
#x1B[1m#x1B[31mE             N Cells:    1#x1B[0m
#x1B[1m#x1B[31mE             N Points:   4#x1B[0m
#x1B[1m#x1B[31mE             X Bounds:   -3.536e-01, 3.536e-01#x1B[0m
#x1B[1m#x1B[31mE             Y Bounds:   -3.536e-01, 3.536e-01#x1B[0m
#x1B[1m#x1B[31mE             Z Bounds:   -3.536e-01, 3.536e-01#x1B[0m
#x1B[1m#x1B[31mE             N Arrays:   3,#x1B[0m
#x1B[1m#x1B[31mE           seed=0,  # or any other generated value#x1B[0m
#x1B[1m#x1B[31mE       )#x1B[0m
#x1B[1m#x1B[31mE       #x1B[0m
#x1B[1m#x1B[31mE       You can reproduce this example by temporarily adding @reproduce_failure('6.150.1', b'AEEA') as a decorator on your test case#x1B[0m

args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64))
fun        = <PjitFunction of <function ffi_call at 0x7f3be7320a40>>
params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
prev       = <object object at 0x7f3ca25a3b00>
prim       = ffi_call

#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/dispatch.py#x1B[0m:91: JaxRuntimeError

tests/warp/energies/elastic/hyperelastic/test_arap_muscle_v2.py::test_arap_muscle_hess_quad

Flake rate in main: 100.00% (Passed 0 times, Failed 15 times)

Stack Traces | 69.8s run time

model = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
mesh = UnstructuredGrid (0x7f3bdf179cc0)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   3

    #x1B[0m#x1B[37m@hypothesis#x1B[39;49;00m.given(seed=testing.seed())#x1B[90m#x1B[39;49;00m
>   #x1B[94mdef#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[92mtest_arap_muscle_hess_quad#x1B[39;49;00m(#x1B[90m#x1B[39;49;00m
                   ^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        seed: #x1B[96mint#x1B[39;49;00m, model: Model, mesh: pv.UnstructuredGrid#x1B[90m#x1B[39;49;00m
    ) -> #x1B[94mNone#x1B[39;49;00m:#x1B[90m#x1B[39;49;00m

f          = <function given.<locals>.run_test_as_given.<locals>.wrapped_test at 0x7f3c401b4e00>
mesh       = UnstructuredGrid (0x7f3bdf179cc0)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   3
model      = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)

#x1B[1m#x1B[.../elastic/hyperelastic/test_arap_muscle_v2.py#x1B[0m:71: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
#x1B[1m#x1B[.../elastic/hyperelastic/test_arap_muscle_v2.py#x1B[0m:74: in test_arap_muscle_hess_quad
    #x1B[0mcommon.check_hess_quad(seed, model, mesh)#x1B[90m#x1B[39;49;00m
        mesh       = UnstructuredGrid (0x7f3bdf179cc0)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   3
        model      = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
        seed       = 0
#x1B[1m#x1B[.../elastic/hyperelastic/common.py#x1B[0m:62: in check_hess_quad
    #x1B[0mexpected: Scalar = jnp.vdot(p, model.hess_prod(u, p))#x1B[90m#x1B[39;49;00m
                                   ^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        actual     = Array(0.5856506, dtype=float64)
        mesh       = UnstructuredGrid (0x7f3bdf179cc0)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   3
        model      = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
        p          = Array([[-0.93533705, -0.11326748,  0.16791879],
       [ 0.81027622, -0.72888169,  0.62500202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)
        seed       = 0
        u          = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
#x1B[1m#x1B[.../apple/model/_model.py#x1B[0m:117: in hess_prod
    #x1B[0moutput_wp: Full = #x1B[96mself#x1B[39;49;00m.warp.hess_prod(u_full, p_full)#x1B[90m#x1B[39;49;00m
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        output_jax = Array([[0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.]], dtype=float64)
        p          = Array([[-0.93533705, -0.11326748,  0.16791879],
       [ 0.81027622, -0.72888169,  0.62500202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)
        p_full     = Array([[-0.93533705, -0.11326748,  0.16791879],
       [ 0.81027622, -0.72888169,  0.62500202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)
        self       = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
        u          = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
        u_full     = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
#x1B[1m#x1B[.../warp/model/_adapter.py#x1B[0m:60: in hess_prod
    #x1B[0m(output,) = #x1B[96mself#x1B[39;49;00m._hess_prod_callable(u, p, output_dims=u.shape)#x1B[90m#x1B[39;49;00m
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        p          = Array([[-0.93533705, -0.11326748,  0.16791879],
       [ 0.81027622, -0.72888169,  0.62500202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)
        self       = WarpModelAdapter(
  wrapped=WarpModel(
    energies={
      'elastic':
      ArapMuscleV2(
        id='elastic',
     ...1,), dtype=float64),
               ),
        clamp_hess_diag=False,
        clamp_hess_quad=False
      )
    }
  )
)
        u          = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
#x1B[1m#x1B[31m.venv/lib/python3.12.../_src/jax_experimental/ffi.py#x1B[0m:640: in __call__
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m call(*args, call_id=call_id)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64))
        call       = <function ffi_call.<locals>.wrapped at 0x7f3becf5b560>
        call_id    = 0
        d          = 0
        device     = 'cpu'
        i          = 1
        input_arg  = <warp._src.jax_experimental.ffi.FfiArg object at 0x7f3bec149e80>
        input_value = Array([[-0.93533705, -0.11326748,  0.16791879],
       [ 0.81027622, -0.72888169,  0.62500202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)
        module     = <warp._src.context.Module object at 0x7f3c205b0fb0>
        num_inputs = 2
        out_types  = [ShapeDtypeStruct(shape=(4, 3), dtype=float64)]
        output_arg = <warp._src.jax_experimental.ffi.FfiArg object at 0x7f3bec14a270>
        output_dims = (4, 3)
        self       = <warp._src.jax_experimental.ffi.FfiCallable object at 0x7f3bec14ade0>
        static_inputs = {}
        vmap_method = 'broadcast_all'
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/ffi.py#x1B[0m:540: in wrapped
    #x1B[0mresults = ffi_call_p.bind(#x1B[90m#x1B[39;49;00m
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64))
        custom_call_api_version = 4
        has_side_effect = False
        in_avals   = [ShapedArray(float64[4,3]), ShapedArray(float64[4,3])]
        input_layouts = None
        input_output_aliases = {}
        kwargs     = {'call_id': 0}
        legacy_backend_config = None
        multiple_results = True
        output_layouts_ = None
        result_avals = (ShapedArray(float64[4,3]),)
        result_shape_dtypes = [ShapeDtypeStruct(shape=(4, 3), dtype=float64)]
        static_input_layouts = ((1, 0), (1, 0))
        static_input_output_aliases = ()
        static_output_layouts = ((1, 0),)
        target_name = 'WarpModelAdapter___hess_prod_callable__locals__hess_prod_callable_316'
        vmap_method = 'broadcast_all'
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:633: in bind
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m #x1B[96mself#x1B[39;49;00m._true_bind(*args, **params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = [Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)]
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        self       = ffi_call
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:649: in _true_bind
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m #x1B[96mself#x1B[39;49;00m.bind_with_trace(prev_trace, args, params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        arg        = Array([[-0.93533705, -0.11326748,  0.16791879],
       [ 0.81027622, -0.72888169,  0.62500202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64))
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        prev_trace = EvalTrace
        self       = ffi_call
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:661: in bind_with_trace
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m trace.process_primitive(#x1B[96mself#x1B[39;49;00m, args, params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64))
        in_type    = [ShapedArray(float64[4,3]), ShapedArray(float64[4,3])]
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        self       = ffi_call
        trace      = EvalTrace
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:1210: in process_primitive
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m primitive.impl(*args, **params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = [Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)]
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        primitive  = ffi_call
        self       = EvalTrace
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

prim = ffi_call
args = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64))
params = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
fun = <PjitFunction of <function ffi_call at 0x7f3becf59da0>>
prev = <object object at 0x7f3ca25a3b00>

    #x1B[0m#x1B[94mdef#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[92mapply_primitive#x1B[39;49;00m(prim, *args, **params):#x1B[90m#x1B[39;49;00m
    #x1B[90m  #x1B[39;49;00m#x1B[33m"""Impl rule that compiles and runs a single primitive 'prim' using XLA."""#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
      fun = xla_primitive_callable(prim, **params)#x1B[90m#x1B[39;49;00m
      #x1B[90m# TODO(yashkatariya): Investigate adding is_primitive to jit and never#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
      #x1B[90m# triggering the disable jit path instead of messing around with it here.#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
      prev = config.disable_jit.swap_local(#x1B[94mFalse#x1B[39;49;00m)#x1B[90m#x1B[39;49;00m
      #x1B[94mtry#x1B[39;49;00m:#x1B[90m#x1B[39;49;00m
>       outs = fun(*args)#x1B[90m#x1B[39;49;00m
               ^^^^^^^^^^#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31mE       jax.errors.JaxRuntimeError: NOT_FOUND: No FFI handler registered for WarpModelAdapter___hess_prod_callable__locals__hess_prod_callable_316 on a platform Host (canonical host)#x1B[0m
#x1B[1m#x1B[31mE       --------------------#x1B[0m
#x1B[1m#x1B[31mE       For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.#x1B[0m
#x1B[1m#x1B[31mE       Falsifying example: test_arap_muscle_hess_quad(#x1B[0m
#x1B[1m#x1B[31mE           model=Model(dirichlet=Dirichlet(dim=3,#x1B[0m
#x1B[1m#x1B[31mE             dirichlet_index=Array([], shape=(0,), dtype=int64),#x1B[0m
#x1B[1m#x1B[31mE             dirichlet_value=Array([], shape=(0,), dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE             free_index=Array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11], dtype=int64),#x1B[0m
#x1B[1m#x1B[31mE             n_points=4),#x1B[0m
#x1B[1m#x1B[31mE            u_full=Array([[-0.76655026, -0.62556086,  0.68767497],#x1B[0m
#x1B[1m#x1B[31mE                   [-0.28233729, -0.69772602, -0.37828513],#x1B[0m
#x1B[1m#x1B[31mE                   [-0.04823423, -1.03855213,  0.95418263],#x1B[0m
#x1B[1m#x1B[31mE                   [ 1.07699311,  0.54753535, -0.98902473]], dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE            jax=JaxModel(energies={}),#x1B[0m
#x1B[1m#x1B[31mE            warp=WarpModelAdapter(wrapped=WarpModel(dim=3,#x1B[0m
#x1B[1m#x1B[31mE              energies={'elastic': ArapMuscleV2(id='elastic',#x1B[0m
#x1B[1m#x1B[31mE                requires_grad=['activation', 'mu'],#x1B[0m
#x1B[1m#x1B[31mE                cells=array(shape=(1,), dtype=vec4i),#x1B[0m
#x1B[1m#x1B[31mE                dhdX=array(shape=(1, 1), dtype=mat43(d)),#x1B[0m
#x1B[1m#x1B[31mE                dV=array(shape=(1, 1), dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE                params=ArapMuscleV2__Params(#x1B[0m
#x1B[1m#x1B[31mE                	activation=array(shape=(1,), dtype=vector(length=6, dtype=float64)),#x1B[0m
#x1B[1m#x1B[31mE                	mu=array(shape=(1,), dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE                ),#x1B[0m
#x1B[1m#x1B[31mE                clamp_hess_diag=False,#x1B[0m
#x1B[1m#x1B[31mE                clamp_hess_quad=False,#x1B[0m
#x1B[1m#x1B[31mE                clamp_lambda=False)})),#x1B[0m
#x1B[1m#x1B[31mE            edges_length_mean=Array(0.99999998, dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE            frozen=False),#x1B[0m
#x1B[1m#x1B[31mE           mesh=UnstructuredGrid (0x7f3bdf179cc0)#x1B[0m
#x1B[1m#x1B[31mE             N Cells:    1#x1B[0m
#x1B[1m#x1B[31mE             N Points:   4#x1B[0m
#x1B[1m#x1B[31mE             X Bounds:   -3.536e-01, 3.536e-01#x1B[0m
#x1B[1m#x1B[31mE             Y Bounds:   -3.536e-01, 3.536e-01#x1B[0m
#x1B[1m#x1B[31mE             Z Bounds:   -3.536e-01, 3.536e-01#x1B[0m
#x1B[1m#x1B[31mE             N Arrays:   3,#x1B[0m
#x1B[1m#x1B[31mE           seed=0,  # or any other generated value#x1B[0m
#x1B[1m#x1B[31mE       )#x1B[0m
#x1B[1m#x1B[31mE       #x1B[0m
#x1B[1m#x1B[31mE       You can reproduce this example by temporarily adding @reproduce_failure('6.150.1', b'AEEA') as a decorator on your test case#x1B[0m

args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64))
fun        = <PjitFunction of <function ffi_call at 0x7f3becf59da0>>
params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
prev       = <object object at 0x7f3ca25a3b00>
prim       = ffi_call

#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/dispatch.py#x1B[0m:91: JaxRuntimeError

tests/warp/energies/elastic/hyperelastic/test_phace.py::test_phace_hess_prod

Flake rate in main: 71.15% (Passed 15 times, Failed 37 times)

Stack Traces | 66.2s run time

model = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
mesh = UnstructuredGrid (0x7f3c22f21cc0)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   5

    #x1B[0m#x1B[37m@hypothesis#x1B[39;49;00m.given(seed=testing.seed())#x1B[90m#x1B[39;49;00m
>   #x1B[94mdef#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[92mtest_phace_hess_prod#x1B[39;49;00m(seed: #x1B[96mint#x1B[39;49;00m, model: Model, mesh: pv.UnstructuredGrid) -> #x1B[94mNone#x1B[39;49;00m:#x1B[90m#x1B[39;49;00m
                   ^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m

f          = <function given.<locals>.run_test_as_given.<locals>.wrapped_test at 0x7f3c401b6de0>
mesh       = UnstructuredGrid (0x7f3c22f21cc0)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   5
model      = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)

#x1B[1m#x1B[.../elastic/hyperelastic/test_phace.py#x1B[0m:62: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
#x1B[1m#x1B[.../elastic/hyperelastic/test_phace.py#x1B[0m:63: in test_phace_hess_prod
    #x1B[0mcommon.check_hess_prod(seed, model, mesh)#x1B[90m#x1B[39;49;00m
        mesh       = UnstructuredGrid (0x7f3c22f21cc0)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   5
        model      = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
        seed       = 0
#x1B[1m#x1B[.../elastic/hyperelastic/common.py#x1B[0m:55: in check_hess_prod
    #x1B[0mtesting.check_jvp(model.grad, model.hess_prod, u, rtol=#x1B[94m1e-3#x1B[39;49;00m)#x1B[90m#x1B[39;49;00m
        mesh       = UnstructuredGrid (0x7f3c22f21cc0)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   5
        model      = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
        seed       = 0
        u          = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
#x1B[1m#x1B[.../apple/model/_model.py#x1B[0m:117: in hess_prod
    #x1B[0moutput_wp: Full = #x1B[96mself#x1B[39;49;00m.warp.hess_prod(u_full, p_full)#x1B[90m#x1B[39;49;00m
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        output_jax = Array([[0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.]], dtype=float64)
        p          = Array([[-0.16308578, -0.56740909,  0.93064292],
       [ 0.14900107,  0.06445298, -0.29018964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64)
        p_full     = Array([[-0.16308578, -0.56740909,  0.93064292],
       [ 0.14900107,  0.06445298, -0.29018964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64)
        self       = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
        u          = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
        u_full     = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
#x1B[1m#x1B[.../warp/model/_adapter.py#x1B[0m:60: in hess_prod
    #x1B[0m(output,) = #x1B[96mself#x1B[39;49;00m._hess_prod_callable(u, p, output_dims=u.shape)#x1B[90m#x1B[39;49;00m
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        p          = Array([[-0.16308578, -0.56740909,  0.93064292],
       [ 0.14900107,  0.06445298, -0.29018964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64)
        self       = WarpModelAdapter(
  wrapped=WarpModel(
    energies={
      'elastic':
      Phace(
        id='elastic',
        requ...1,), dtype=float64),
               ),
        clamp_hess_diag=False,
        clamp_hess_quad=False
      )
    }
  )
)
        u          = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
#x1B[1m#x1B[31m.venv/lib/python3.12.../_src/jax_experimental/ffi.py#x1B[0m:640: in __call__
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m call(*args, call_id=call_id)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64))
        call       = <function ffi_call.<locals>.wrapped at 0x7f3bed377a60>
        call_id    = 0
        d          = 0
        device     = 'cpu'
        i          = 1
        input_arg  = <warp._src.jax_experimental.ffi.FfiArg object at 0x7f3be4b72ea0>
        input_value = Array([[-0.16308578, -0.56740909,  0.93064292],
       [ 0.14900107,  0.06445298, -0.29018964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64)
        module     = <warp._src.context.Module object at 0x7f3c205b0fb0>
        num_inputs = 2
        out_types  = [ShapeDtypeStruct(shape=(4, 3), dtype=float64)]
        output_arg = <warp._src.jax_experimental.ffi.FfiArg object at 0x7f3be4b734d0>
        output_dims = (4, 3)
        self       = <warp._src.jax_experimental.ffi.FfiCallable object at 0x7f3be4b72090>
        static_inputs = {}
        vmap_method = 'broadcast_all'
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/ffi.py#x1B[0m:540: in wrapped
    #x1B[0mresults = ffi_call_p.bind(#x1B[90m#x1B[39;49;00m
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64))
        custom_call_api_version = 4
        has_side_effect = False
        in_avals   = [ShapedArray(float64[4,3]), ShapedArray(float64[4,3])]
        input_layouts = None
        input_output_aliases = {}
        kwargs     = {'call_id': 0}
        legacy_backend_config = None
        multiple_results = True
        output_layouts_ = None
        result_avals = (ShapedArray(float64[4,3]),)
        result_shape_dtypes = [ShapeDtypeStruct(shape=(4, 3), dtype=float64)]
        static_input_layouts = ((1, 0), (1, 0))
        static_input_output_aliases = ()
        static_output_layouts = ((1, 0),)
        target_name = 'WarpModelAdapter___hess_prod_callable__locals__hess_prod_callable_103'
        vmap_method = 'broadcast_all'
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:633: in bind
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m #x1B[96mself#x1B[39;49;00m._true_bind(*args, **params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = [Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64)]
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        self       = ffi_call
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:649: in _true_bind
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m #x1B[96mself#x1B[39;49;00m.bind_with_trace(prev_trace, args, params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        arg        = Array([[-0.16308578, -0.56740909,  0.93064292],
       [ 0.14900107,  0.06445298, -0.29018964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64)
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64))
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        prev_trace = EvalTrace
        self       = ffi_call
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:661: in bind_with_trace
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m trace.process_primitive(#x1B[96mself#x1B[39;49;00m, args, params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64))
        in_type    = [ShapedArray(float64[4,3]), ShapedArray(float64[4,3])]
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        self       = ffi_call
        trace      = EvalTrace
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:1210: in process_primitive
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m primitive.impl(*args, **params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = [Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64)]
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        primitive  = ffi_call
        self       = EvalTrace
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

prim = ffi_call
args = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64))
params = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
fun = <PjitFunction of <function ffi_call at 0x7f3bed376ca0>>
prev = <object object at 0x7f3ca25a3b00>

    #x1B[0m#x1B[94mdef#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[92mapply_primitive#x1B[39;49;00m(prim, *args, **params):#x1B[90m#x1B[39;49;00m
    #x1B[90m  #x1B[39;49;00m#x1B[33m"""Impl rule that compiles and runs a single primitive 'prim' using XLA."""#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
      fun = xla_primitive_callable(prim, **params)#x1B[90m#x1B[39;49;00m
      #x1B[90m# TODO(yashkatariya): Investigate adding is_primitive to jit and never#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
      #x1B[90m# triggering the disable jit path instead of messing around with it here.#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
      prev = config.disable_jit.swap_local(#x1B[94mFalse#x1B[39;49;00m)#x1B[90m#x1B[39;49;00m
      #x1B[94mtry#x1B[39;49;00m:#x1B[90m#x1B[39;49;00m
>       outs = fun(*args)#x1B[90m#x1B[39;49;00m
               ^^^^^^^^^^#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31mE       jax.errors.JaxRuntimeError: NOT_FOUND: No FFI handler registered for WarpModelAdapter___hess_prod_callable__locals__hess_prod_callable_103 on a platform Host (canonical host)#x1B[0m
#x1B[1m#x1B[31mE       --------------------#x1B[0m
#x1B[1m#x1B[31mE       For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.#x1B[0m
#x1B[1m#x1B[31mE       Falsifying example: test_phace_hess_prod(#x1B[0m
#x1B[1m#x1B[31mE           model=Model(dirichlet=Dirichlet(dim=3,#x1B[0m
#x1B[1m#x1B[31mE             dirichlet_index=Array([], shape=(0,), dtype=int64),#x1B[0m
#x1B[1m#x1B[31mE             dirichlet_value=Array([], shape=(0,), dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE             free_index=Array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11], dtype=int64),#x1B[0m
#x1B[1m#x1B[31mE             n_points=4),#x1B[0m
#x1B[1m#x1B[31mE            u_full=Array([[-0.42033094, -0.31086302, -1.18779184],#x1B[0m
#x1B[1m#x1B[31mE                   [ 1.08350666, -0.84188598,  0.03346199],#x1B[0m
#x1B[1m#x1B[31mE                   [ 0.06173197, -0.09646734, -0.07578768],#x1B[0m
#x1B[1m#x1B[31mE                   [-0.78159042, -1.11343283, -1.00513326]], dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE            jax=JaxModel(energies={}),#x1B[0m
#x1B[1m#x1B[31mE            warp=WarpModelAdapter(wrapped=WarpModel(dim=3,#x1B[0m
#x1B[1m#x1B[31mE              energies={'elastic': Phace(id='elastic',#x1B[0m
#x1B[1m#x1B[31mE                requires_grad=['activation', 'lambda_', 'mu', 'muscle_fraction'],#x1B[0m
#x1B[1m#x1B[31mE                cells=array(shape=(1,), dtype=vec4i),#x1B[0m
#x1B[1m#x1B[31mE                dhdX=array(shape=(1, 1), dtype=mat43(d)),#x1B[0m
#x1B[1m#x1B[31mE                dV=array(shape=(1, 1), dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE                params=Phace__Params(#x1B[0m
#x1B[1m#x1B[31mE                	activation=array(shape=(1,), dtype=vector(length=6, dtype=float64)),#x1B[0m
#x1B[1m#x1B[31mE                	lambda_=array(shape=(1,), dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE                	mu=array(shape=(1,), dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE                	muscle_fraction=array(shape=(1,), dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE                ),#x1B[0m
#x1B[1m#x1B[31mE                clamp_hess_diag=False,#x1B[0m
#x1B[1m#x1B[31mE                clamp_hess_quad=False,#x1B[0m
#x1B[1m#x1B[31mE                clamp_lambda=False)})),#x1B[0m
#x1B[1m#x1B[31mE            edges_length_mean=Array(0.99999998, dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE            frozen=False),#x1B[0m
#x1B[1m#x1B[31mE           mesh=UnstructuredGrid (0x7f3c22f21cc0)#x1B[0m
#x1B[1m#x1B[31mE             N Cells:    1#x1B[0m
#x1B[1m#x1B[31mE             N Points:   4#x1B[0m
#x1B[1m#x1B[31mE             X Bounds:   -3.536e-01, 3.536e-01#x1B[0m
#x1B[1m#x1B[31mE             Y Bounds:   -3.536e-01, 3.536e-01#x1B[0m
#x1B[1m#x1B[31mE             Z Bounds:   -3.536e-01, 3.536e-01#x1B[0m
#x1B[1m#x1B[31mE             N Arrays:   5,#x1B[0m
#x1B[1m#x1B[31mE           seed=0,  # or any other generated value#x1B[0m
#x1B[1m#x1B[31mE       )#x1B[0m
#x1B[1m#x1B[31mE       #x1B[0m
#x1B[1m#x1B[31mE       You can reproduce this example by temporarily adding @reproduce_failure('6.150.1', b'AEEA') as a decorator on your test case#x1B[0m

args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...8964],
       [ 0.76601795,  0.26512078,  0.06760015],
       [-0.61669224,  0.7083533 , -0.88112142]], dtype=float64))
fun        = <PjitFunction of <function ffi_call at 0x7f3bed376ca0>>
params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
prev       = <object object at 0x7f3ca25a3b00>
prim       = ffi_call

#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/dispatch.py#x1B[0m:91: JaxRuntimeError

tests/warp/energies/elastic/hyperelastic/test_phace.py::test_phace_hess_quad

Flake rate in main: 71.15% (Passed 15 times, Failed 37 times)

Stack Traces | 70.7s run time

model = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
mesh = UnstructuredGrid (0x7f3c22f21cc0)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   5

    #x1B[0m#x1B[37m@hypothesis#x1B[39;49;00m.given(seed=testing.seed())#x1B[90m#x1B[39;49;00m
>   #x1B[94mdef#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[92mtest_phace_hess_quad#x1B[39;49;00m(seed: #x1B[96mint#x1B[39;49;00m, model: Model, mesh: pv.UnstructuredGrid) -> #x1B[94mNone#x1B[39;49;00m:#x1B[90m#x1B[39;49;00m
                   ^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m

f          = <function given.<locals>.run_test_as_given.<locals>.wrapped_test at 0x7f3c401b7240>
mesh       = UnstructuredGrid (0x7f3c22f21cc0)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   5
model      = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)

#x1B[1m#x1B[.../elastic/hyperelastic/test_phace.py#x1B[0m:67: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 
#x1B[1m#x1B[.../elastic/hyperelastic/test_phace.py#x1B[0m:68: in test_phace_hess_quad
    #x1B[0mcommon.check_hess_quad(seed, model, mesh)#x1B[90m#x1B[39;49;00m
        mesh       = UnstructuredGrid (0x7f3c22f21cc0)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   5
        model      = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
        seed       = 0
#x1B[1m#x1B[.../elastic/hyperelastic/common.py#x1B[0m:62: in check_hess_quad
    #x1B[0mexpected: Scalar = jnp.vdot(p, model.hess_prod(u, p))#x1B[90m#x1B[39;49;00m
                                   ^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        actual     = Array(0.63912849, dtype=float64)
        mesh       = UnstructuredGrid (0x7f3c22f21cc0)
  N Cells:    1
  N Points:   4
  X Bounds:   -3.536e-01, 3.536e-01
  Y Bounds:   -3.536e-01, 3.536e-01
  Z Bounds:   -3.536e-01, 3.536e-01
  N Arrays:   5
        model      = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
        p          = Array([[-0.93533705, -0.11326748,  0.16791879],
       [ 0.81027622, -0.72888169,  0.62500202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)
        seed       = 0
        u          = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
#x1B[1m#x1B[.../apple/model/_model.py#x1B[0m:117: in hess_prod
    #x1B[0moutput_wp: Full = #x1B[96mself#x1B[39;49;00m.warp.hess_prod(u_full, p_full)#x1B[90m#x1B[39;49;00m
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        output_jax = Array([[0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.],
       [0., 0., 0.]], dtype=float64)
        p          = Array([[-0.93533705, -0.11326748,  0.16791879],
       [ 0.81027622, -0.72888169,  0.62500202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)
        p_full     = Array([[-0.93533705, -0.11326748,  0.16791879],
       [ 0.81027622, -0.72888169,  0.62500202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)
        self       = Model(
  dirichlet=Dirichlet(
    dim=3,
    dirichlet_index=Array([], shape=(0,), dtype=int64),
    dirichlet_value=A...se,
          clamp_hess_quad=False
        )
      }
    )
  ),
  edges_length_mean=Array(0.99999998, dtype=float64)
)
        u          = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
        u_full     = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
#x1B[1m#x1B[.../warp/model/_adapter.py#x1B[0m:60: in hess_prod
    #x1B[0m(output,) = #x1B[96mself#x1B[39;49;00m._hess_prod_callable(u, p, output_dims=u.shape)#x1B[90m#x1B[39;49;00m
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        p          = Array([[-0.93533705, -0.11326748,  0.16791879],
       [ 0.81027622, -0.72888169,  0.62500202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)
        self       = WarpModelAdapter(
  wrapped=WarpModel(
    energies={
      'elastic':
      Phace(
        id='elastic',
        requ...1,), dtype=float64),
               ),
        clamp_hess_diag=False,
        clamp_hess_quad=False
      )
    }
  )
)
        u          = Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656,  0.32470532,  0.08279294],
       [-0.75529065,  0.86755208, -1.07914894]], dtype=float64)
#x1B[1m#x1B[31m.venv/lib/python3.12.../_src/jax_experimental/ffi.py#x1B[0m:640: in __call__
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m call(*args, call_id=call_id)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64))
        call       = <function ffi_call.<locals>.wrapped at 0x7f3bdf8db4c0>
        call_id    = 0
        d          = 0
        device     = 'cpu'
        i          = 1
        input_arg  = <warp._src.jax_experimental.ffi.FfiArg object at 0x7f3c2072a180>
        input_value = Array([[-0.93533705, -0.11326748,  0.16791879],
       [ 0.81027622, -0.72888169,  0.62500202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)
        module     = <warp._src.context.Module object at 0x7f3c205b0fb0>
        num_inputs = 2
        out_types  = [ShapeDtypeStruct(shape=(4, 3), dtype=float64)]
        output_arg = <warp._src.jax_experimental.ffi.FfiArg object at 0x7f3c2072a660>
        output_dims = (4, 3)
        self       = <warp._src.jax_experimental.ffi.FfiCallable object at 0x7f3c205b10a0>
        static_inputs = {}
        vmap_method = 'broadcast_all'
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/ffi.py#x1B[0m:540: in wrapped
    #x1B[0mresults = ffi_call_p.bind(#x1B[90m#x1B[39;49;00m
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64))
        custom_call_api_version = 4
        has_side_effect = False
        in_avals   = [ShapedArray(float64[4,3]), ShapedArray(float64[4,3])]
        input_layouts = None
        input_output_aliases = {}
        kwargs     = {'call_id': 0}
        legacy_backend_config = None
        multiple_results = True
        output_layouts_ = None
        result_avals = (ShapedArray(float64[4,3]),)
        result_shape_dtypes = [ShapeDtypeStruct(shape=(4, 3), dtype=float64)]
        static_input_layouts = ((1, 0), (1, 0))
        static_input_output_aliases = ()
        static_output_layouts = ((1, 0),)
        target_name = 'WarpModelAdapter___hess_prod_callable__locals__hess_prod_callable_210'
        vmap_method = 'broadcast_all'
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:633: in bind
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m #x1B[96mself#x1B[39;49;00m._true_bind(*args, **params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = [Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)]
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        self       = ffi_call
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:649: in _true_bind
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m #x1B[96mself#x1B[39;49;00m.bind_with_trace(prev_trace, args, params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        arg        = Array([[-0.93533705, -0.11326748,  0.16791879],
       [ 0.81027622, -0.72888169,  0.62500202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64))
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        prev_trace = EvalTrace
        self       = ffi_call
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:661: in bind_with_trace
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m trace.process_primitive(#x1B[96mself#x1B[39;49;00m, args, params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64))
        in_type    = [ShapedArray(float64[4,3]), ShapedArray(float64[4,3])]
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        self       = ffi_call
        trace      = EvalTrace
#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/core.py#x1B[0m:1210: in process_primitive
    #x1B[0m#x1B[94mreturn#x1B[39;49;00m primitive.impl(*args, **params)#x1B[90m#x1B[39;49;00m
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^#x1B[90m#x1B[39;49;00m
        args       = [Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64)]
        params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
        primitive  = ffi_call
        self       = EvalTrace
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ 

prim = ffi_call
args = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64))
params = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
fun = <PjitFunction of <function ffi_call at 0x7f3bdf8daa20>>
prev = <object object at 0x7f3ca25a3b00>

    #x1B[0m#x1B[94mdef#x1B[39;49;00m#x1B[90m #x1B[39;49;00m#x1B[92mapply_primitive#x1B[39;49;00m(prim, *args, **params):#x1B[90m#x1B[39;49;00m
    #x1B[90m  #x1B[39;49;00m#x1B[33m"""Impl rule that compiles and runs a single primitive 'prim' using XLA."""#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
      fun = xla_primitive_callable(prim, **params)#x1B[90m#x1B[39;49;00m
      #x1B[90m# TODO(yashkatariya): Investigate adding is_primitive to jit and never#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
      #x1B[90m# triggering the disable jit path instead of messing around with it here.#x1B[39;49;00m#x1B[90m#x1B[39;49;00m
      prev = config.disable_jit.swap_local(#x1B[94mFalse#x1B[39;49;00m)#x1B[90m#x1B[39;49;00m
      #x1B[94mtry#x1B[39;49;00m:#x1B[90m#x1B[39;49;00m
>       outs = fun(*args)#x1B[90m#x1B[39;49;00m
               ^^^^^^^^^^#x1B[90m#x1B[39;49;00m
#x1B[1m#x1B[31mE       jax.errors.JaxRuntimeError: NOT_FOUND: No FFI handler registered for WarpModelAdapter___hess_prod_callable__locals__hess_prod_callable_210 on a platform Host (canonical host)#x1B[0m
#x1B[1m#x1B[31mE       --------------------#x1B[0m
#x1B[1m#x1B[31mE       For simplicity, JAX has removed its internal frames from the traceback of the following exception. Set JAX_TRACEBACK_FILTERING=off to include these.#x1B[0m
#x1B[1m#x1B[31mE       Falsifying example: test_phace_hess_quad(#x1B[0m
#x1B[1m#x1B[31mE           model=Model(dirichlet=Dirichlet(dim=3,#x1B[0m
#x1B[1m#x1B[31mE             dirichlet_index=Array([], shape=(0,), dtype=int64),#x1B[0m
#x1B[1m#x1B[31mE             dirichlet_value=Array([], shape=(0,), dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE             free_index=Array([ 0,  1,  2,  3,  4,  5,  6,  7,  8,  9, 10, 11], dtype=int64),#x1B[0m
#x1B[1m#x1B[31mE             n_points=4),#x1B[0m
#x1B[1m#x1B[31mE            u_full=Array([[-0.81671112, -0.09798451,  1.16489133],#x1B[0m
#x1B[1m#x1B[31mE                   [ 1.09169105,  1.08964695,  0.47490621],#x1B[0m
#x1B[1m#x1B[31mE                   [ 0.39828523,  0.91929163, -0.3556751 ],#x1B[0m
#x1B[1m#x1B[31mE                   [-1.03554067,  0.96252309, -0.89989933]], dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE            jax=JaxModel(energies={}),#x1B[0m
#x1B[1m#x1B[31mE            warp=WarpModelAdapter(wrapped=WarpModel(dim=3,#x1B[0m
#x1B[1m#x1B[31mE              energies={'elastic': Phace(id='elastic',#x1B[0m
#x1B[1m#x1B[31mE                requires_grad=['activation', 'lambda_', 'mu', 'muscle_fraction'],#x1B[0m
#x1B[1m#x1B[31mE                cells=array(shape=(1,), dtype=vec4i),#x1B[0m
#x1B[1m#x1B[31mE                dhdX=array(shape=(1, 1), dtype=mat43(d)),#x1B[0m
#x1B[1m#x1B[31mE                dV=array(shape=(1, 1), dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE                params=Phace__Params(#x1B[0m
#x1B[1m#x1B[31mE                	activation=array(shape=(1,), dtype=vector(length=6, dtype=float64)),#x1B[0m
#x1B[1m#x1B[31mE                	lambda_=array(shape=(1,), dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE                	mu=array(shape=(1,), dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE                	muscle_fraction=array(shape=(1,), dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE                ),#x1B[0m
#x1B[1m#x1B[31mE                clamp_hess_diag=False,#x1B[0m
#x1B[1m#x1B[31mE                clamp_hess_quad=False,#x1B[0m
#x1B[1m#x1B[31mE                clamp_lambda=False)})),#x1B[0m
#x1B[1m#x1B[31mE            edges_length_mean=Array(0.99999998, dtype=float64),#x1B[0m
#x1B[1m#x1B[31mE            frozen=False),#x1B[0m
#x1B[1m#x1B[31mE           mesh=UnstructuredGrid (0x7f3c22f21cc0)#x1B[0m
#x1B[1m#x1B[31mE             N Cells:    1#x1B[0m
#x1B[1m#x1B[31mE             N Points:   4#x1B[0m
#x1B[1m#x1B[31mE             X Bounds:   -3.536e-01, 3.536e-01#x1B[0m
#x1B[1m#x1B[31mE             Y Bounds:   -3.536e-01, 3.536e-01#x1B[0m
#x1B[1m#x1B[31mE             Z Bounds:   -3.536e-01, 3.536e-01#x1B[0m
#x1B[1m#x1B[31mE             N Arrays:   5,#x1B[0m
#x1B[1m#x1B[31mE           seed=0,  # or any other generated value#x1B[0m
#x1B[1m#x1B[31mE       )#x1B[0m
#x1B[1m#x1B[31mE       #x1B[0m
#x1B[1m#x1B[31mE       You can reproduce this example by temporarily adding @reproduce_failure('6.150.1', b'AEEA') as a decorator on your test case#x1B[0m

args       = (Array([[-0.19973847, -0.69493137,  1.13980015],
       [ 0.18248829,  0.07893845, -0.35540827],
       [ 0.93817656, ...0202],
       [ 0.07366043, -0.68028063, -0.04946978],
       [-1.00891584, -0.82771482, -0.90290814]], dtype=float64))
fun        = <PjitFunction of <function ffi_call at 0x7f3bdf8daa20>>
params     = {'attributes': (('call_id', 0),), 'custom_call_api_version': 4, 'has_side_effect': False, 'input_layouts': ((1, 0), (1, 0)), ...}
prev       = <object object at 0x7f3ca25a3b00>
prim       = ffi_call

#x1B[1m#x1B[31m.venv/lib/python3.12.../jax/_src/dispatch.py#x1B[0m:91: JaxRuntimeError

To view more test analytics, go to the Test Analytics Dashboard
_{📋 Got 3 mins? Take this short survey to help us improve Test Analytics.}

fix(deps): update dependency jax to >=0.9,<0.10

9813468

renovate bot assigned liblaf Jan 21, 2026

renovate bot requested a review from liblaf January 21, 2026 00:29

renovate-approve bot previously approved these changes Jan 21, 2026

View reviewed changes

renovate bot temporarily deployed to MegaLinter January 21, 2026 00:29 Inactive

chore(ci): auto fixes from pre-commit.com hooks

3c7c646

for more information, see https://pre-commit.ci

pre-commit-ci bot dismissed renovate-approve[bot]’s stale review via 3c7c646 January 21, 2026 00:29

pre-commit-ci bot deployed to MegaLinter January 21, 2026 00:29 Active

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Uh oh!

fix(deps): update dependency jax to >=0.9,<0.10 #84

fix(deps): update dependency jax to >=0.9,<0.10 #84

Uh oh!

renovate bot commented Jan 21, 2026 •

edited

Loading

Uh oh!

renovate bot commented Jan 21, 2026

Uh oh!

liblaf bot commented Jan 21, 2026 •

edited

Loading

Uh oh!

codecov bot commented Jan 21, 2026 •

edited

Loading

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Uh oh!

fix(deps): update dependency jax to >=0.9,<0.10 #84

Are you sure you want to change the base?

fix(deps): update dependency jax to >=0.9,<0.10 #84

Uh oh!

Conversation

renovate bot commented Jan 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Release Notes

v0.9.0

Configuration

Uh oh!

renovate bot commented Jan 21, 2026

⚠️ Artifact update problem

File name: uv.lock

Uh oh!

liblaf bot commented Jan 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

✅⚠️MegaLinter analysis: Success with warnings

Detailed Issues

Uh oh!

codecov bot commented Jan 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

❌ 12 Tests Failed:

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

renovate bot commented Jan 21, 2026 •

edited

Loading

`v0.9.0`

liblaf bot commented Jan 21, 2026 •

edited

Loading

codecov bot commented Jan 21, 2026 •

edited

Loading