ENH: _restore_axes accomodates for negative axis arguments

andgoldschmidt · Jun 5, 2024 · 8cde00d · 8cde00d
1 parent 9474e7f
commit 8cde00d
Showing 1 changed file with 4 additions and 1 deletion.
diff --git a/derivative/differentiation.py b/derivative/differentiation.py
@@ -248,7 +248,10 @@ def _restore_axes(dX: NDArray, axis: int, orig_shape: tuple[int, ...]) -> NDArra
         return dX.flatten()
     else:
         # order of operations coupled with _align_axes
-        extra_dims = tuple(length for ax, length in enumerate(orig_shape) if ax != axis)
+        orig_diff_axis = range(len(orig_shape))[axis] # to handle negative axis args
+        extra_dims = tuple(
+            length for ax, length in enumerate(orig_shape) if ax != orig_diff_axis
+        )
         moved_shape = (orig_shape[axis],) + extra_dims
         dX = np.moveaxis(dX.T.reshape((moved_shape)), 0, axis)
         return dX