linting: pyright & mypy

Author: cakedev0

src/array_api_extra/_delegation.py

@@ -1110,26 +1110,26 @@ def quantile(
     )
     device = get_device(a)
     a = xp.asarray(a, dtype=dtype, device=device)
-    q = xp.asarray(q, dtype=dtype, device=device)
+    q_arr = xp.asarray(q, dtype=dtype, device=device)
     # TODO: cast weights here? Assert weights are on the same device as `a`?
 
-    if xp.any((q > 1) | (q < 0) | xp.isnan(q)):
+    if xp.any((q_arr > 1) | (q_arr < 0) | xp.isnan(q_arr)):
         msg = "`q` values must be in the range [0, 1]"
         raise ValueError(msg)
 
     # Delegate where possible.
     if is_numpy_namespace(xp) and nan_policy == "propagate":
-        return xp.quantile(a, q, axis=axis, method=method, keepdims=keepdims, weights=weights)
+        return xp.quantile(a, q_arr, axis=axis, method=method, keepdims=keepdims, weights=weights)
     # No delegation for dask: I couldn't make it work
     basic_case = method == "linear" and weights is None and nan_policy == "propagate"
     if (basic_case and is_jax_namespace(xp)) or is_cupy_namespace(xp):
-        return xp.quantile(a, q, axis=axis, method=method, keepdims=keepdims)
+        return xp.quantile(a, q_arr, axis=axis, method=method, keepdims=keepdims)
     if basic_case and is_torch_namespace(xp):
-        return xp.quantile(a, q, dim=axis, interpolation=method, keepdim=keepdims)
+        return xp.quantile(a, q_arr, dim=axis, interpolation=method, keepdim=keepdims)
 
     # XXX: I'm not sure we want to support dask, it seems uterly slow...
     # Otherwise call our implementation (will sort data)
     return _quantile.quantile(
-        a, q, axis=axis, method=method, keepdims=keepdims,
+        a, q_arr, axis=axis, method=method, keepdims=keepdims,
         nan_policy=nan_policy, weights=weights, xp=xp
     )

src/array_api_extra/_lib/_quantile.py

@@ -48,10 +48,10 @@ def quantile(  # numpydoc ignore=PR01,RT01
         if not axis_none:
             res = xp.moveaxis(res, axis, 0)
     else:
-        weights = xp.asarray(weights, dtype=xp.float64, device=device)
+        weights_arr = xp.asarray(weights, dtype=xp.float64, device=device)
         average = method == 'averaged_inverted_cdf'
         res = _weighted_quantile(
-            a, q, weights, n, axis, average,
+            a, q, weights_arr, n, axis, average,
             nan_policy=nan_policy, xp=xp, device=device
         )
 
@@ -80,7 +80,7 @@ def _quantile(  # numpydoc ignore=GL08
     if method == "linear":
         m = 1 - q
     else: # method is "inverted_cdf" or "averaged_inverted_cdf"
-        m = 0
+        m = xp.asarray(0, dtype=q.dtype)
 
     jg = q * float(n) + m - 1
 
@@ -112,7 +112,6 @@ def _weighted_quantile(
     """
     a is expected to be 1d or 2d.
     """
-    kwargs = dict(n=n, average=average, nan_policy=nan_policy, xp=xp, device=device)
     a = xp.moveaxis(a, axis, -1)
     if weights.ndim > 1:
         weights = xp.moveaxis(weights, axis, -1)
@@ -121,15 +120,15 @@ def _weighted_quantile(
     if a.ndim == 1:
         x = xp.take(a, sorter)
         w = xp.take(weights, sorter)
-        return _weighted_quantile_sorted_1d(x, q, w, **kwargs)
+        return _weighted_quantile_sorted_1d(x, q, w, n, average, nan_policy, xp, device)
 
     d, = eager_shape(a, axis=0)
     res = []
     for idx in range(d):
         w = weights if weights.ndim == 1 else weights[idx, ...]
         w = xp.take(w, sorter[idx, ...])
         x = xp.take(a[idx, ...], sorter[idx, ...])
-        res.append(_weighted_quantile_sorted_1d(x, q, w, **kwargs))
+        res.append(_weighted_quantile_sorted_1d(x, q, w, n, average, nan_policy, xp, device))
     res = xp.stack(res, axis=1)
     return res
 

tests/test_funcs.py

@@ -1578,18 +1578,18 @@ def test_weighted_against_numpy(self, xp: ModuleType, keepdims: bool, q_np: Arra
         rng = np.random.default_rng()
         n, d = 10, 20
         a_np = rng.random((n, d))
-        kwargs = dict(keepdims=keepdims)
         mask_nan = np.zeros((n, d), dtype=bool)
-        if nan_policy != "no_nans":
+        if nan_policy == "no_nans":
+            nan_policy = "propagate"
+        else:
             # from 0% to 100% of NaNs:
             mask_nan = rng.random((n, d)) < rng.random((n, 1))
             # don't put nans in the first row:
             mask_nan[:] = False
             a_np[mask_nan] = np.nan
-            kwargs['nan_policy'] = nan_policy
 
-        a = xp.asarray(a_np)
-        q = xp.asarray(np.copy(q_np))
+        a = xp.asarray(a_np, copy=True)
+        q = xp.asarray(q_np, copy=True)
         m = 'inverted_cdf'
 
         np_quantile = np.quantile
@@ -1604,23 +1604,24 @@ def test_weighted_against_numpy(self, xp: ModuleType, keepdims: bool, q_np: Arra
             (rng.integers(0, 2, (n, d)), 0),
             (rng.integers(0, 2, (n, d)), 1),
         ]:
-            print(w_np)
             with warnings.catch_warnings(record=True) as warning:
                 warnings.filterwarnings("always", "invalid value encountered in divide", RuntimeWarning)
                 warnings.filterwarnings("ignore", "All-NaN slice encountered", RuntimeWarning)
                 try:
-                    expected = np_quantile(a_np, q_np, axis=axis, method=m, weights=w_np, keepdims=keepdims)
+                    expected = np_quantile(  # type: ignore[call-overload]
+                        a_np, np.asarray(q_np),
+                        axis=axis, method=m, weights=w_np, keepdims=keepdims
+                    )
                 except IndexError:
-                    print('index error')
                     continue
                 if warning:  # this means some weights sum was 0, in this case we skip calling xpx.quantile
-                    print('warning')
                     continue
             expected = xp.asarray(expected)
-            print("not skiped")
 
             w = xp.asarray(w_np)
-            actual = quantile(a, q, axis=axis, method=m, weights=w, **kwargs)
+            actual = quantile(  
+                a, q, axis=axis, method=m, weights=w, keepdims=keepdims, nan_policy=nan_policy
+            )
             xp_assert_close(actual, expected, atol=1e-12)
 
     def test_2d_axis(self, xp: ModuleType):