jvp should only return derivative

zmheiko · zmheiko · commit 3eaaa8eac30b · 2025-04-23T11:24:47.000+02:00
diff --git a/demo/cfd/cfd-tesseract/tesseract_api.py b/demo/cfd/cfd-tesseract/tesseract_api.py
@@ -204,7 +204,7 @@ def jvp_jit(
         filtered_apply,
         [flatten_with_paths(inputs, include_paths=jvp_inputs)],
         [tangent_vector],
-    )
+    )[1]
 
 
 @eqx.filter_jit
