diff --git a/hj_reachability/finite_differences/upwind_first.py b/hj_reachability/finite_differences/upwind_first.py
index 6c0c9fa..892cdc2 100644
--- a/hj_reachability/finite_differences/upwind_first.py
+++ b/hj_reachability/finite_differences/upwind_first.py
@@ -166,7 +166,7 @@ def _diff_coefficients(k: Optional[int] = None, stencil: Optional[Array] = None)
                              f"{(k, stencil.shape[-1] - 1)}.")
     return np.linalg.solve(
         np.diff(poly.polyvander(stencil, k), axis=-2)[..., 1:].swapaxes(-1, -2),
-        np.eye(k)[(np.newaxis,) * (stencil.ndim - 1) + (0,)])
+        np.eye(k)[(np.newaxis,) * (stencil.ndim - 1) + (0, ..., np.newaxis)])[..., 0]
 
 
 def _substencil_coefficients(k: int) -> Array: