black

Author: crusaderky

array_api_compat/common/_helpers.py

@@ -788,6 +788,7 @@ def to_device(x: Array, device: Device, /, *, stream: Optional[Union[int, Any]]
         return x
     return x.to_device(device, stream=stream)
 
+
 def size(x):
     """
     Return the total number of elements of x.
@@ -802,6 +803,7 @@ def size(x):
         return None
     return math.prod(x.shape)
 
+
 def is_writeable_array(x):
     """
     Return False if x.__setitem__ is expected to raise; True otherwise
@@ -812,6 +814,7 @@ def is_writeable_array(x):
         return False
     return True
 
+
 def _parse_copy_param(x, copy: bool | None | Literal["_force_false"]) -> bool:
     """Preprocess and validate a copy parameter, in line with the same
     parameter in np.asarray(), np.astype(), etc.
@@ -827,8 +830,10 @@ def _parse_copy_param(x, copy: bool | None | Literal["_force_false"]) -> bool:
         raise ValueError(f"Invalid value for copy: {copy!r}")
     return copy
 
+
 _undef = object()
 
+
 class at:
     """
     Update operations for read-only arrays.
@@ -897,6 +902,7 @@ class at:
     --------
     https://jax.readthedocs.io/en/latest/_autosummary/jax.numpy.ndarray.at.html
     """
+
     __slots__ = ("x", "idx")
 
     def __init__(self, x, idx=_undef):
@@ -945,7 +951,7 @@ def _common(
         if copy and is_jax_array(self.x):
             # Use JAX's at[]
             at_ = self.x.at[self.idx]
-            args = (y, ) if y is not _undef else ()
+            args = (y,) if y is not _undef else ()
             return getattr(at_, at_op)(*args, **kwargs), None
 
         # Emulate at[] behaviour for non-JAX arrays
@@ -958,12 +964,9 @@ def get(self, copy: bool | None = True, **kwargs):
         # Special case when xp=numpy and idx is a fancy index
         # If copy is not False, avoid an unnecessary double copy.
         # if copy is forced to False, raise.
-        if (
-            is_numpy_array(self.x) 
-            and (
-                isinstance(self.idx, (list, tuple))
-                or (is_numpy_array(self.idx) and self.idx.dtype.kind in "biu")
-            )
+        if is_numpy_array(self.x) and (
+            isinstance(self.idx, (list, tuple))
+            or (is_numpy_array(self.idx) and self.idx.dtype.kind in "biu")
         ):
             if copy is False:
                 raise ValueError(
@@ -994,12 +997,14 @@ def apply(self, ufunc, /, **kwargs):
         ufunc.at(x, self.idx)
         return x
 
-    def _iop(self, at_op: str, elwise_op: Callable[[Array, Array], Array], y: Array, **kwargs):
+    def _iop(
+        self, at_op: str, elwise_op: Callable[[Array, Array], Array], y: Array, **kwargs
+    ):
         """x[idx] += y or equivalent in-place operation on a subset of x
 
         which is the same as saying
             x[idx] = x[idx] + y
-        Note that this is not the same as 
+        Note that this is not the same as
             operator.iadd(x[idx], y)
         Consider for example when x is a numpy array and idx is a fancy index, which
         triggers a deep copy on __getitem__.
@@ -1017,11 +1022,11 @@ def add(self, y, /, **kwargs):
     def subtract(self, y, /, **kwargs):
         """x[idx] -= y"""
         return self._iop("subtract", operator.sub, y, **kwargs)
-    
+
     def multiply(self, y, /, **kwargs):
         """x[idx] *= y"""
         return self._iop("multiply", operator.mul, y, **kwargs)
-    
+
     def divide(self, y, /, **kwargs):
         """x[idx] /= y"""
         return self._iop("divide", operator.truediv, y, **kwargs)
@@ -1040,9 +1045,10 @@ def max(self, y, /, **kwargs):
         xp = array_namespace(self.x)
         return self._iop("max", xp.maximum, y, **kwargs)
 
+
 def where(condition, x=None, y=None, /, copy: bool | None = True):
     """Return elements from x when condition is True and from y when
-    it is False. 
+    it is False.
 
     This is a wrapper around xp.where that adds the copy parameter: