MAINT: move normalization logic to _normalizations

Author: ev-br

torch_np/_binary_ufuncs.py

@@ -1,19 +1,14 @@
+from . import _helpers
 from ._detail import _binary_ufuncs
+from ._normalizations import ArrayLike, DTypeLike, SubokLike, normalizer
 
 __all__ = [
     name for name in dir(_binary_ufuncs) if not name.startswith("_") and name != "torch"
 ]
 
 
-from . import _helpers
-from ._detail import _util
-
-# TODO: consolidate normalizations
-from ._funcs import ArrayLike, DTypeLike, SubokLike, normalizer
-
-
 def deco_binary_ufunc(torch_func):
-    """Common infra for unary ufuncs.
+    """Common infra for binary ufuncs.
 
     Normalize arguments, sort out type casting, broadcasting and delegate to
     the pytorch functions for the actual work.
@@ -33,19 +28,9 @@ def wrapped(
         signature=None,
         extobj=None,
     ):
-        if order != "K" or not where or signature or extobj:
-            raise NotImplementedError
-
-        # XXX: dtype=... parameter
-        if dtype is not None:
-            raise NotImplementedError
-
-        out_shape_dtype = None
-        if out is not None:
-            out_shape_dtype = (out.get().dtype, out.get().shape)
-
-        tensors = _util.cast_and_broadcast((x1, x2), out_shape_dtype, casting)
-
+        tensors = _helpers.ufunc_preprocess(
+            (x1, x2), out, where, casting, order, dtype, subok, signature, extobj
+        )
         result = torch_func(*tensors)
         return _helpers.result_or_out(result, out)
 

torch_np/_decorators.py

@@ -5,8 +5,6 @@
 from . import _dtypes, _helpers
 from ._detail import _util
 
-NoValue = None
-
 
 def out_shape_dtype(func):
     """Handle out=... kwarg for ufuncs.

torch_np/_funcs.py

@@ -1,159 +1,18 @@
-import operator
-import typing
-from typing import Optional, Sequence
+from typing import Optional
 
 import torch
 
-from . import _decorators, _helpers
-from ._detail import _dtypes_impl, _flips, _reductions, _util
+from . import _helpers
+from ._detail import _flips, _reductions, _util
 from ._detail import implementations as _impl
-
-################################## normalizations
-
-ArrayLike = typing.TypeVar("ArrayLike")
-DTypeLike = typing.TypeVar("DTypeLike")
-SubokLike = typing.TypeVar("SubokLike")
-AxisLike = typing.TypeVar("AxisLike")
-
-# annotate e.g. atleast_1d(*arys)
-UnpackedSeqArrayLike = typing.TypeVar("UnpackedSeqArrayLike")
-
-
-import inspect
-
-from . import _dtypes
-
-
-def normalize_array_like(x, name=None):
-    (tensor,) = _helpers.to_tensors(x)
-    return tensor
-
-
-def normalize_optional_array_like(x, name=None):
-    # This explicit normalizer is needed because otherwise normalize_array_like
-    # does not run for a parameter annotated as Optional[ArrayLike]
-    return None if x is None else normalize_array_like(x, name)
-
-
-def normalize_seq_array_like(x, name=None):
-    tensors = _helpers.to_tensors(*x)
-    return tensors
-
-
-def normalize_dtype(dtype, name=None):
-    # cf _decorators.dtype_to_torch
-    torch_dtype = None
-    if dtype is not None:
-        dtype = _dtypes.dtype(dtype)
-        torch_dtype = dtype.torch_dtype
-    return torch_dtype
-
-
-def normalize_subok_like(arg, name):
-    if arg:
-        raise ValueError(f"'{name}' parameter is not supported.")
-
-
-def normalize_axis_like(arg, name=None):
-    from ._ndarray import ndarray
-
-    if isinstance(arg, ndarray):
-        arg = operator.index(arg)
-    return arg
-
-
-normalizers = {
-    ArrayLike: normalize_array_like,
-    Optional[ArrayLike]: normalize_optional_array_like,
-    Sequence[ArrayLike]: normalize_seq_array_like,
-    UnpackedSeqArrayLike: normalize_seq_array_like,  # cf handling in normalize
-    DTypeLike: normalize_dtype,
-    SubokLike: normalize_subok_like,
-    AxisLike: normalize_axis_like,
-}
-
-import functools
-
-
-def normalize_this(arg, parm):
-    """Normalize arg if a normalizer is registred."""
-    normalizer = normalizers.get(parm.annotation, None)
-    if normalizer:
-        return normalizer(arg)
-    else:
-        # untyped arguments pass through
-        return arg
-
-
-def normalizer(func):
-    @functools.wraps(func)
-    def wrapped(*args, **kwds):
-        sig = inspect.signature(func)
-
-        # first, check for *args in positional parameters. Case in point:
-        # atleast_1d(*arys: UnpackedSequenceArrayLike)
-        # if found,  consume all args into a tuple to normalize as a whole
-        for j, param in enumerate(sig.parameters.values()):
-            if param.annotation == UnpackedSeqArrayLike:
-                if j == 0:
-                    args = (args,)
-                else:
-                    # args = args[:j] + (args[j:],) would likely work
-                    # not present in numpy codebase, so do not bother just yet.
-                    # NB: branching on j ==0 is to avoid the empty tuple, args[:j]
-                    raise NotImplementedError
-                break
-
-        # loop over positional parameters and actual arguments
-        lst, dct = [], {}
-        for arg, (name, parm) in zip(args, sig.parameters.items()):
-            print(arg, name, parm.annotation)
-            lst.append(normalize_this(arg, parm))
-
-        # normalize keyword arguments
-        for name, arg in kwds.items():
-            if not name in sig.parameters:
-                # unknown kwarg, bail out
-                raise TypeError(
-                    f"{func.__name__}() got an unexpected keyword argument '{name}'."
-                )
-
-            print("kw: ", name, sig.parameters[name].annotation)
-            parm = sig.parameters[name]
-            dct[name] = normalize_this(arg, parm)
-
-        ba = sig.bind(*lst, **dct)
-        ba.apply_defaults()
-
-        # Now that all parameters have been consumed, check:
-        # Anything that has not been bound is unexpected positional arg => raise.
-        # If there are too few actual arguments, this fill fail in func(*ba.args) below
-        if len(args) > len(ba.args):
-            raise TypeError(
-                f"{func.__name__}() takes {len(ba.args)} positional argument but {len(args)} were given."
-            )
-
-        # TODO:
-        # 1. [LOOKS OK] kw-only parameters : see vstack
-        # 2. [LOOKS OK] extra unknown args -- error out : nonzero([2, 0, 3], oops=42)
-        # 3. [LOOKS OK] optional (tensor_or_none) : untyped => pass through
-        # 4. [LOOKS OK] DTypeLike : positional or kw
-        # 5. axes : live in _impl or in types? several ways of handling them
-        # 6. [OK, NOT HERE] keepdims : peel off, postprocess
-        # 7. OutLike : normal & keyword-only, peel off, postprocess
-        # 8. [LOOKS OK] *args
-        # 9. consolidate normalizations (_funcs, _wrapper)
-        # 10. consolidate decorators (_{unary,binary}_ufuncs)
-        # 11. out= arg : validate it's an ndarray
-
-        # finally, pass normalized arguments through
-        result = func(*ba.args, **ba.kwargs)
-        return result
-
-    return wrapped
-
-
-##################################
+from ._normalizations import (
+    ArrayLike,
+    AxisLike,
+    DTypeLike,
+    SubokLike,
+    UnpackedSeqArrayLike,
+    normalizer,
+)
 
 
 @normalizer

torch_np/_helpers.py

@@ -41,6 +41,26 @@ def cast_and_broadcast(tensors, out, casting):
     return tuple(tensors)
 
 
+def ufunc_preprocess(
+    tensors, out, where, casting, order, dtype, subok, signature, extobj
+):
+    # internal preprocessing or args in ufuncs (cf _unary_ufuncs, _binary_ufuncs)
+    if order != "K" or not where or signature or extobj:
+        raise NotImplementedError
+
+    # XXX: dtype=... parameter
+    if dtype is not None:
+        raise NotImplementedError
+
+    out_shape_dtype = None
+    if out is not None:
+        out_shape_dtype = (out.get().dtype, out.get().shape)
+
+    tensors = _util.cast_and_broadcast(tensors, out_shape_dtype, casting)
+
+    return tensors
+
+
 # ### Return helpers: wrap a single tensor, a tuple of tensors, out= etc ###