Make enzyme_ref delete in_shapes instead of out_shapes (#32)

EnzymeAD · Jan 31, 2024 · 185b4c0 · 185b4c0
1 parent 2115612
commit 185b4c0
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/src/enzyme_ad/jax/primitives.py b/src/enzyme_ad/jax/primitives.py
@@ -284,7 +284,7 @@ def _enzyme_rev_impl(
     lang: enzyme_call.Language,
     pipeline_options
 ) -> Sequence[jax.Array]:
-    del args_flat, source, out_shapes
+    del args_flat, source, in_shapes
     raise RuntimeError("must be JIT'ed")