feat(numpy): support negative indices in checked array accessors

include/pybind11/numpy.h

@@ -1156,12 +1156,14 @@ class array : public buffer {
     }
 #endif
 
-    /// Dimension along a given axis
+    /// Dimension along a given axis. Negative axes count from the end.
+    /// Throws if dim is outside [-ndim(), ndim()).
     ssize_t shape(ssize_t dim) const {
-        if (dim >= ndim()) {
+        const ssize_t axis = dim < 0 ? dim + ndim() : dim;
+        if (static_cast<size_t>(axis) >= static_cast<size_t>(ndim())) {
             fail_dim_check(dim, "invalid axis");
         }
-        return shape()[dim];
+        return shape()[axis];
     }
 
     /// Strides of the array
@@ -1174,12 +1176,14 @@ class array : public buffer {
     }
 #endif
 
-    /// Stride along a given axis
+    /// Stride along a given axis. Negative axes count from the end.
+    /// Throws if dim is outside [-ndim(), ndim()).
     ssize_t strides(ssize_t dim) const {
-        if (dim >= ndim()) {
+        const ssize_t axis = dim < 0 ? dim + ndim() : dim;
+        if (static_cast<size_t>(axis) >= static_cast<size_t>(ndim())) {
             fail_dim_check(dim, "invalid axis");
         }
-        return strides()[dim];
+        return strides()[axis];
     }
 
     /// Return the NumPy array flags
@@ -1196,23 +1200,24 @@ class array : public buffer {
     }
 
     /// Pointer to the contained data. If index is not provided, points to the
-    /// beginning of the buffer. May throw if the index would lead to out of bounds access.
+    /// beginning of the buffer. Negative indices count from the end of each dimension.
+    /// Throws if an index is out of bounds.
     template <typename... Ix>
     const void *data(Ix... index) const {
         return static_cast<const void *>(detail::array_proxy(m_ptr)->data + offset_at(index...));
     }
 
     /// Mutable pointer to the contained data. If index is not provided, points to the
-    /// beginning of the buffer. May throw if the index would lead to out of bounds access.
-    /// May throw if the array is not writeable.
+    /// beginning of the buffer. Negative indices count from the end of each dimension.
+    /// Throws if an index is out of bounds or the array is not writeable.
     template <typename... Ix>
     void *mutable_data(Ix... index) {
         check_writeable();
         return static_cast<void *>(detail::array_proxy(m_ptr)->data + offset_at(index...));
     }
 
     /// Byte offset from beginning of the array to a given index (full or partial).
-    /// May throw if the index would lead to out of bounds access.
+    /// Negative indices count from the end of each dimension. Throws if an index is out of bounds.
     template <typename... Ix>
     ssize_t offset_at(Ix... index) const {
         if ((ssize_t) sizeof...(index) > ndim()) {
@@ -1224,7 +1229,7 @@ class array : public buffer {
     ssize_t offset_at() const { return 0; }
 
     /// Item count from beginning of the array to a given index (full or partial).
-    /// May throw if the index would lead to out of bounds access.
+    /// Negative indices count from the end of each dimension. Throws if an index is out of bounds.
     template <typename... Ix>
     ssize_t index_at(Ix... index) const {
         return offset_at(index...) / itemsize();
@@ -1337,8 +1342,7 @@ class array : public buffer {
 
     template <typename... Ix>
     ssize_t byte_offset(Ix... index) const {
-        check_dimensions(index...);
-        return detail::byte_offset_unsafe(strides(), ssize_t(index)...);
+        return byte_offset_impl(ssize_t(0), shape(), strides(), ssize_t(index)...);
     }
 
     void check_writeable() const {
@@ -1347,21 +1351,20 @@ class array : public buffer {
         }
     }
 
-    template <typename... Ix>
-    void check_dimensions(Ix... index) const {
-        check_dimensions_impl(ssize_t(0), shape(), ssize_t(index)...);
-    }
-
-    void check_dimensions_impl(ssize_t, const ssize_t *) const {}
+    ssize_t byte_offset_impl(ssize_t, const ssize_t *, const ssize_t *) const { return 0; }
 
     template <typename... Ix>
-    void check_dimensions_impl(ssize_t axis, const ssize_t *shape, ssize_t i, Ix... index) const {
-        if (i >= *shape) {
+    ssize_t byte_offset_impl(
+        ssize_t axis, const ssize_t *shape, const ssize_t *strides, ssize_t i, Ix... index) const {
+        const ssize_t normalized = i < 0 ? i + *shape : i;
+        // Unsigned comparison also rejects indices that remain negative after normalization.
+        if (static_cast<size_t>(normalized) >= static_cast<size_t>(*shape)) {
             throw index_error(std::string("index ") + std::to_string(i)
                               + " is out of bounds for axis " + std::to_string(axis)
                               + " with size " + std::to_string(*shape));
         }
-        check_dimensions_impl(axis + 1, shape + 1, index...);
+        const ssize_t remaining = byte_offset_impl(axis + 1, shape + 1, strides + 1, index...);
+        return normalized * *strides + remaining;
     }
 
     /// Create array from any object -- always returns a new reference
@@ -1449,7 +1452,8 @@ class array_t : public array {
         return static_cast<T *>(array::mutable_data(index...));
     }
 
-    // Reference to element at a given index
+    // Reference to element at a given index. Negative indices count from the end of each
+    // dimension. Throws if out of bounds.
     template <typename... Ix>
     const T &at(Ix... index) const {
         if ((ssize_t) sizeof...(index) != ndim()) {
@@ -1459,7 +1463,8 @@ class array_t : public array {
                  + byte_offset(ssize_t(index)...) / itemsize());
     }
 
-    // Mutable reference to element at a given index
+    // Mutable reference to element at a given index. Negative indices count from the end of each
+    // dimension. Throws if out of bounds.
     template <typename... Ix>
     T &mutable_at(Ix... index) {
         if ((ssize_t) sizeof...(index) != ndim()) {

tests/test_numpy_array.cpp

@@ -135,9 +135,10 @@ arr_t &mutate_at_t(arr_t &a, Ix... idx) {
 
 #define def_index_fn(name, type)                                                                  \
     sm.def(#name, [](type a) { return name(a); });                                                \
-    sm.def(#name, [](type a, int i) { return name(a, i); });                                      \
-    sm.def(#name, [](type a, int i, int j) { return name(a, i, j); });                            \
-    sm.def(#name, [](type a, int i, int j, int k) { return name(a, i, j, k); });
+    sm.def(#name, [](type a, py::ssize_t i) { return name(a, i); });                              \
+    sm.def(#name, [](type a, py::ssize_t i, py::ssize_t j) { return name(a, i, j); });            \
+    sm.def(#name,                                                                                 \
+           [](type a, py::ssize_t i, py::ssize_t j, py::ssize_t k) { return name(a, i, j, k); });
 
 template <typename T, typename T2>
 py::handle auxiliaries(T &&r, T2 &&r2) {

tests/test_numpy_array.py

@@ -69,6 +69,10 @@ def test_array_attributes():
     with pytest.raises(IndexError) as excinfo:
         m.strides(a, 2)
     assert str(excinfo.value) == "invalid axis: 2 (ndim = 2)"
+    assert m.shape(a, -1) == 3
+    assert m.shape(a, -2) == 2
+    assert m.strides(a, -1) == 2
+    assert m.strides(a, -2) == 6
     assert not m.writeable(a)
     assert m.size(a) == 6
     assert m.itemsize(a) == 2
@@ -204,6 +208,8 @@ def test_bounds_check(arr):
     for func in (
         m.index_at,
         m.index_at_t,
+        m.offset_at,
+        m.offset_at_t,
         m.data,
         m.data_t,
         m.mutate_data,
@@ -217,6 +223,110 @@ def test_bounds_check(arr):
         with pytest.raises(IndexError) as excinfo:
             func(arr, 0, 4)
         assert str(excinfo.value) == "index 4 is out of bounds for axis 1 with size 3"
+        with pytest.raises(IndexError) as excinfo:
+            func(arr, -3, 0)
+        assert str(excinfo.value) == "index -3 is out of bounds for axis 0 with size 2"
+        with pytest.raises(IndexError) as excinfo:
+            func(arr, 0, -4)
+        assert str(excinfo.value) == "index -4 is out of bounds for axis 1 with size 3"
+
+
+@pytest.mark.parametrize("shape", [(), (0,), (2, 0), (2, 3)])
+def test_negative_axes(shape):
+    a = np.empty(shape, dtype="=u2")
+    for func, values in ((m.shape, a.shape), (m.strides, a.strides)):
+        for axis in range(-a.ndim, a.ndim):
+            assert func(a, axis) == values[axis]
+        for axis in (-sys.maxsize - 1, -a.ndim - 1, a.ndim, sys.maxsize):
+            with pytest.raises(IndexError) as excinfo:
+                func(a, axis)
+            assert str(excinfo.value) == f"invalid axis: {axis} (ndim = {a.ndim})"
+
+
+@pytest.mark.parametrize(
+    "indices", [(-1,), (-2,), (-1, 0), (0, -1), (-2, 0), (0, -3), (-2, -3), (-1, -1)]
+)
+def test_negative_indices(arr, indices):
+    normalized = tuple(i + n if i < 0 else i for i, n in zip(indices, arr.shape))
+    funcs = (
+        m.index_at,
+        m.index_at_t,
+        m.offset_at,
+        m.offset_at_t,
+        m.data,
+        m.data_t,
+        m.mutate_data,
+        m.mutate_data_t,
+    )
+    if len(indices) == arr.ndim:
+        funcs += (m.at_t, m.mutate_at_t)
+    for func in funcs:
+        actual = arr.copy()
+        expected = arr.copy()
+        np.testing.assert_array_equal(
+            func(actual, *indices), func(expected, *normalized)
+        )
+        np.testing.assert_array_equal(actual, expected)
+
+
+@pytest.mark.parametrize("shape", [(0,), (2, 0), (2, 3)])
+def test_negative_index_boundaries(shape):
+    a = np.zeros(shape, dtype="=u2")
+    for func in (
+        m.index_at,
+        m.index_at_t,
+        m.offset_at,
+        m.offset_at_t,
+        m.data,
+        m.data_t,
+        m.mutate_data,
+        m.mutate_data_t,
+        m.at_t,
+        m.mutate_at_t,
+    ):
+        for axis, size in enumerate(shape):
+            for value in (-sys.maxsize - 1, -size - 1, size, sys.maxsize):
+                indices = [0] * len(shape)
+                indices[axis] = value
+                bad_axis = next(
+                    dim
+                    for dim, (i, n) in enumerate(zip(indices, shape))
+                    if i < -n or i >= n
+                )
+                with pytest.raises(IndexError) as excinfo:
+                    func(a, *indices)
+                assert str(excinfo.value) == (
+                    f"index {indices[bad_axis]} is out of bounds for axis {bad_axis} "
+                    f"with size {shape[bad_axis]}"
+                )
+
+
+@pytest.mark.parametrize("layout", ["C", "F", "reversed", "broadcast"])
+def test_negative_indices_strided(layout):
+    a = np.arange(6, dtype="=u2").reshape(2, 3)
+    if layout == "F":
+        a = np.asfortranarray(a)
+    elif layout == "reversed":
+        a = a[::-1, ::-1]
+    elif layout == "broadcast":
+        a = np.broadcast_to(a[:1, :], (2, 3))
+    for i in range(-2, 2):
+        for j in range(-3, 3):
+            assert m.at_t(a, i, j) == a[i, j]
+            offset = (i + 2 if i < 0 else i) * a.strides[0]
+            offset += (j + 3 if j < 0 else j) * a.strides[1]
+            assert m.offset_at(a, i, j) == offset
+            assert m.offset_at_t(a, i, j) == offset
+            assert m.index_at(a, i, j) == offset // a.itemsize
+            assert m.index_at_t(a, i, j) == offset // a.itemsize
+            if a.flags.writeable:
+                expected = int(a[i, j]) + 1
+                m.mutate_at_t(a, i, j)
+                assert a[i, j] == expected
+    a.flags.writeable = False
+    for func in (m.mutate_data, m.mutate_data_t, m.mutate_at_t):
+        with pytest.raises(ValueError, match="array is not writeable"):
+            func(a, -1, -1)
 
 
 def test_make_c_f_array():