Merge branch 'main' into pad-delegate

Author: lucascolley

pixi.lock

@@ -204,12 +204,13 @@ files = ["src", "tests"]
 python_version = "3.10"
 warn_unused_configs = true
 strict = true
-enable_error_code = ["ignore-without-code", "redundant-expr", "truthy-bool"]
-warn_unreachable = true
+enable_error_code = ["ignore-without-code", "truthy-bool"]
 disallow_untyped_defs = false
 disallow_incomplete_defs = false
 # data-apis/array-api#589
 disallow_any_expr = false
+# false positives with input validation
+disable_error_code = ["redundant-expr", "unreachable"]
 
 [[tool.mypy.overrides]]
 module = "array_api_extra.*"
@@ -232,6 +233,8 @@ reportExplicitAny = false
 reportUnknownMemberType = false
 # no array-api-compat type stubs
 reportUnknownVariableType = false
+# false positives for input validation
+reportUnreachable = false
 
 
 # Ruff

pyproject.toml

@@ -73,7 +73,7 @@ def _delegate(xp: ModuleType, *backends: IsNamespace) -> bool:
 
 def pad(
     x: Array,
-    pad_width: int,
+    pad_width: int | tuple[int, int] | list[tuple[int, int]],
     mode: str = "constant",
     *,
     constant_values: bool | int | float | complex = 0,
@@ -86,8 +86,12 @@ def pad(
     ----------
     x : array
         Input array.
-    pad_width : int
+    pad_width : int or tuple of ints or list of pairs of ints
         Pad the input array with this many elements from each side.
+        If a list of tuples, ``[(before_0, after_0), ... (before_N, after_N)]``,
+        each pair applies to the corresponding axis of ``x``.
+        A single tuple, ``(before, after)``, is equivalent to a list of ``x.ndim``
+        copies of this tuple.
     mode : str, optional
         Only "constant" mode is currently supported, which pads with
         the value passed to `constant_values`.

src/array_api_extra/_delegation.py

@@ -551,19 +551,47 @@ def sinc(x: Array, /, *, xp: ModuleType | None = None) -> Array:
 
 def pad(
     x: Array,
-    pad_width: int,
+    pad_width: int | tuple[int, int] | list[tuple[int, int]],
     *,
     constant_values: bool | int | float | complex = 0,
     xp: ModuleType,
 ) -> Array:  # numpydoc ignore=PR01,RT01
     """See docstring in `array_api_extra._delegation.py`."""
+    # make pad_width a list of length-2 tuples of ints
+    x_ndim = cast(int, x.ndim)
+    if isinstance(pad_width, int):
+        pad_width = [(pad_width, pad_width)] * x_ndim
+    if isinstance(pad_width, tuple):
+        pad_width = [pad_width] * x_ndim
+
+    # https://github.com/data-apis/array-api-extra/pull/82#discussion_r1905688819
+    slices: list[slice] = []  # type: ignore[no-any-explicit]
+    newshape: list[int] = []
+    for ax, w_tpl in enumerate(pad_width):
+        if len(w_tpl) != 2:
+            msg = f"expect a 2-tuple (before, after), got {w_tpl}."
+            raise ValueError(msg)
+
+        sh = x.shape[ax]
+        if w_tpl[0] == 0 and w_tpl[1] == 0:
+            sl = slice(None, None, None)
+        else:
+            start, stop = w_tpl
+            stop = None if stop == 0 else -stop
+
+            sl = slice(start, stop, None)
+            sh += w_tpl[0] + w_tpl[1]
+
+        newshape.append(sh)
+        slices.append(sl)
+
     padded = xp.full(
-        tuple(x + 2 * pad_width for x in x.shape),
+        tuple(newshape),
         fill_value=constant_values,
         dtype=x.dtype,
         device=_compat.device(x),
     )
-    padded[(slice(pad_width, -pad_width, None),) * x.ndim] = x
+    padded[tuple(slices)] = x
     return padded
 
 
@@ -613,22 +641,39 @@ class at:  # pylint: disable=invalid-name  # numpydoc ignore=PR02
 
     Warnings
     --------
-    (a) When you omit the ``copy`` parameter, you should always immediately overwrite
-    the parameter array::
+    (a) When you omit the ``copy`` parameter, you should never reuse the parameter
+    array later on; ideally, you should reassign it immediately::
 
         >>> import array_api_extra as xpx
         >>> x = xpx.at(x, 0).set(2)
 
-    The anti-pattern below must be avoided, as it will result in different
-    behaviour on read-only versus writeable arrays::
+    The above best practice pattern ensures that the behaviour won't change depending
+    on whether ``x`` is writeable or not, as the original ``x`` object is dereferenced
+    as soon as ``xpx.at`` returns; this way there is no risk to accidentally update it
+    twice.
+
+    On the reverse, the anti-pattern below must be avoided, as it will result in
+    different behaviour on read-only versus writeable arrays::
 
         >>> x = xp.asarray([0, 0, 0])
         >>> y = xpx.at(x, 0).set(2)
         >>> z = xpx.at(x, 1).set(3)
 
-    In the above example, ``x == [0, 0, 0]``, ``y == [2, 0, 0]`` and z == ``[0, 3, 0]``
-    when ``x`` is read-only, whereas ``x == y == z == [2, 3, 0]`` when ``x`` is
-    writeable!
+    In the above example, both calls to ``xpx.at`` update ``x`` in place *if possible*.
+    This causes the behaviour to diverge depending on whether ``x`` is writeable or not:
+
+    - If ``x`` is writeable, then after the snippet above you'll have
+      ``x == y == z == [2, 3, 0]``
+    - If ``x`` is read-only, then you'll end up with
+      ``x == [0, 0, 0]``, ``y == [2, 0, 0]`` and ``z == [0, 3, 0]``.
+
+    The correct pattern to use if you want diverging outputs from the same input is
+    to enforce copies::
+
+        >>> x = xp.asarray([0, 0, 0])
+        >>> y = xpx.at(x, 0).set(2, copy=True)  # Never updates x
+        >>> z = xpx.at(x, 1).set(3)  # May or may not update x in place
+        >>> del x  # avoid accidental reuse of x as we don't know its state anymore
 
     (b) The array API standard does not support integer array indices.
     The behaviour of update methods when the index is an array of integers is
@@ -728,7 +773,7 @@ def _update_common(
             raise ValueError(msg)
 
         if copy not in (True, False, None):
-            msg = f"copy must be True, False, or None; got {copy!r}"  # pyright: ignore[reportUnreachable]
+            msg = f"copy must be True, False, or None; got {copy!r}"
             raise ValueError(msg)
 
         if copy is None:

src/array_api_extra/_lib/_funcs.py

@@ -416,3 +416,22 @@ def test_device(self):
 
     def test_xp(self):
         assert_array_equal(pad(xp.asarray(0), 1, xp=xp), xp.zeros(3))
+
+    def test_tuple_width(self):
+        a = xp.reshape(xp.arange(12), (3, 4))
+        padded = pad(a, (1, 0))
+        assert padded.shape == (4, 5)
+
+        padded = pad(a, (1, 2))
+        assert padded.shape == (6, 7)
+
+        with pytest.raises(ValueError, match="expect a 2-tuple"):
+            pad(a, [(1, 2, 3)])  # type: ignore[list-item]  # pyright: ignore[reportArgumentType]
+
+    def test_list_of_tuples_width(self):
+        a = xp.reshape(xp.arange(12), (3, 4))
+        padded = pad(a, [(1, 0), (0, 2)])
+        assert padded.shape == (4, 6)
+
+        padded = pad(a, [(1, 0), (0, 0)])
+        assert padded.shape == (4, 4)