Remove the kwargs from jnp.reshape for compatibility with older JAX v…

…ersions (#752)
pnkraemer · Jun 13, 2024 · ede2d82 · ede2d82
1 parent ef9a6b6
commit ede2d82
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/probdiffeq/backend/numpy.py b/probdiffeq/backend/numpy.py
@@ -49,7 +49,7 @@ def diff_along_axis(arr, /, *, axis):
 
 
 def reshape(arr, /, new_shape, order="C"):
-    return jnp.reshape(arr, shape=new_shape, order=order)
+    return jnp.reshape(arr, new_shape, order=order)
 
 
 def flip(arr, /):