Apply suggestions from code review

Co-authored-by: Lucas Colley <lucas.colley8@gmail.com>

Author: bruAristimunha

src/array_api_extra/_delegation.py

@@ -264,7 +264,7 @@ def diag_indices(
     Returns
     -------
     tuple of array
-        ``ndim`` 1-D integer arrays of length ``n`` that together index
+        1-D integer arrays of length ``n`` that together index
         the main diagonal of an array of shape ``(n,) * ndim``.
 
     Examples

src/array_api_extra/_lib/_funcs.py

@@ -350,7 +350,7 @@ def create_diagonal(
 
 
 def diag_indices(
-    n: int, /, *, ndim: int = 2, device: Device | None = None, xp: ModuleType
+    n: int, /, *, ndim: int, device: Device | None, xp: ModuleType
 ) -> tuple[Array, ...]:  # numpydoc ignore=PR01,RT01
     """See docstring in array_api_extra._delegation."""
     idx = xp.arange(n, device=device)
@@ -368,8 +368,8 @@ def _tri_indices(
 ) -> tuple[Array, Array]:  # numpydoc ignore=PR01,RT01
     """Shared implementation for `tril_indices` and `triu_indices`."""
     cols = n if m is None else m
-    rows = xp.arange(n, device=device)[:, None]
-    cols_a = xp.arange(cols, device=device)[None, :]
+    rows = xp.arange(n, device=device)[:, xp.newaxis]
+    cols_a = xp.arange(cols, device=device)[xp.newaxis, :]
     delta = cols_a - rows
     mask = delta >= offset if upper else delta <= offset
     r, c = xp.nonzero(mask)
@@ -380,9 +380,9 @@ def tril_indices(
     n: int,
     /,
     *,
-    offset: int = 0,
-    m: int | None = None,
-    device: Device | None = None,
+    offset: int,
+    m: int | None,
+    device: Device | None,
     xp: ModuleType,
 ) -> tuple[Array, Array]:  # numpydoc ignore=PR01,RT01
     """See docstring in array_api_extra._delegation."""
@@ -393,9 +393,9 @@ def triu_indices(
     n: int,
     /,
     *,
-    offset: int = 0,
-    m: int | None = None,
-    device: Device | None = None,
+    offset: int,
+    m: int | None,
+    device: Device | None,
     xp: ModuleType,
 ) -> tuple[Array, Array]:  # numpydoc ignore=PR01,RT01
     """See docstring in array_api_extra._delegation."""

tests/test_funcs.py

@@ -812,7 +812,7 @@ def test_torch(self, torch: ModuleType):
 @pytest.mark.xfail_xp_backend(Backend.SPARSE, reason="no arange", strict=False)
 class TestDiagIndices:
     def test_basic(self, xp: ModuleType):
-        rows, cols = diag_indices(5, xp=xp)
+        rows, cols = diag_indices(5)
         ref_rows, ref_cols = np.diag_indices(5)
         xp_assert_equal(rows, xp.asarray(ref_rows))
         xp_assert_equal(cols, xp.asarray(ref_cols))
@@ -865,7 +865,7 @@ def test_basic(
         xpx_fn: Callable[..., tuple[Array, Array]],
         np_fn: Callable[..., tuple[Array, Array]],
     ):
-        rows, cols = xpx_fn(4, xp=xp)
+        rows, cols = xpx_fn(4)
         ref_rows, ref_cols = np_fn(4)
         xp_assert_equal(rows, xp.asarray(ref_rows))
         xp_assert_equal(cols, xp.asarray(ref_cols))