Merge pull request #363 from rsokl/ufunc-overload

Implements Tensor.__array_ufunc__

Author: rsokl

src/mygrad/tensor_base.py

@@ -6,6 +6,7 @@
 
 from numbers import Integral, Number
 from typing import (
+    TYPE_CHECKING,
     Any,
     Callable,
     Dict,
@@ -66,6 +67,10 @@
 
 __all__ = ["Tensor", "asarray", "astensor"]
 
+if TYPE_CHECKING:  # pragma: no cover
+    from mygrad.ufuncs._ufunc_creators import ufunc as mygrad_ufunc
+
+
 CONSTANT_ONLY_DTYPES = (np.integer, np.bool_)
 
 
@@ -349,6 +354,53 @@ def astensor(
     return tensor(t, dtype=dtype, constant=constant, copy=False, ndmin=0)
 
 
+_REGISTERED_UFUNC: Dict[np.ufunc, Type["mygrad_ufunc"]] = {}
+_REGISTERED_BOOL_ONLY_UFUNC: Set[np.ufunc] = {
+    np.isnan,
+    np.isfinite,
+    np.isinf,
+    np.isnat,
+    np.signbit,
+    np.logical_not,
+    np.logical_and,
+    np.logical_or,
+    np.logical_xor,
+    np.greater,
+    np.greater_equal,
+    np.less,
+    np.less_equal,
+    np.equal,
+}
+
+# These are ufuncs that users might mistake for being differentiable functions;
+# for this reason we make explicit the fact that only constant tensors are permitted
+# in these operations.
+_REGISTERED_CONST_ONLY_UFUNC = {
+    np.floor_divide,
+    np.remainder,
+    np.mod,
+    np.fmod,
+    np.divmod,
+    np.rint,
+    np.sign,
+    np.floor,
+    np.ceil,
+    np.trunc,
+}
+
+
+class _ConstantOnly(ValueError):
+    pass
+
+
+def _as_constant_array(t: Union["Tensor", np.ndarray]) -> np.ndarray:
+    if isinstance(t, Tensor):
+        if t.constant is False:
+            raise _ConstantOnly()
+        return t.data
+    return t
+
+
 class Tensor:
     """A numpy-array-like object capable of serving as a node in a computational
     graph that supports back-propagation of derivatives via the chain rule.
@@ -505,6 +557,97 @@ class Tensor:
 
     __array_priority__ = 15.0
 
+    def __array_ufunc__(
+        self, ufunc: Type[np.ufunc], method: str, *inputs: ArrayLike, **kwargs
+    ) -> Union["Tensor", np.ndarray]:
+        """An interface provided by NumPy to override the behavior of its ufuncs [1]_.
+
+        MyGrad implements its own ufuncs for all differentiable NumPy ufuncs.
+
+        Non-differentiable numpy ufuncs simply get called on the underlying arrays of tensors and
+        will return ndarrays.
+
+        The differentiability - or lack thereof - of ufuncs may not be obvious to end users.
+        Thus potentially ambiguous ufuncs (e.g. `numpy.ceil`) will be made to raise on non-constant
+        tensors so that the lack of differentiability is made obvious to the users. This design decision
+        is made in the same spirit as requiring integer-dtype tensors be constant.
+
+        References
+        ----------
+        .. [1] https://numpy.org/doc/stable/reference/arrays.classes.html#numpy.class.__array_ufunc__
+
+        Examples
+        --------
+        NumPy ufuncs that represent differentiable operations are overloaded by MyGrad tensors
+        so that they support backprop
+
+        >>> import mygrad as mg
+        >>> import numpy as np
+
+        >>> x = mg.tensor([1., 2.])
+
+        This calls ``mygrad.sin`` under the hood.
+
+        >>> np.sin(x)  # returns a tensor
+        Tensor([0.84147098, 0.90929743])
+
+        >>> np.sin(x).backward()
+        >>> x.grad  # note: derivative of
+        array([ 0.54030231, -0.41614684])
+
+        Specifying a dtype, a ``where`` mask, an in-place target (via ``out``) as an array
+        or a tensor, are all supported.
+
+        >>> x = mg.tensor([1., 2.])
+        >>> y = mg.tensor([-1., -1.])
+        >>> np.exp(x, where=[False, True], out=y)
+        Tensor([-1.       ,  7.3890561])
+        >>> y.backward()
+        >>> x.grad
+        array([0.       , 7.3890561])
+
+        Non-differentiable NumPy ufuncs simply operate on the ndarrays that are wrapped
+        by MyGrad tensors; these return ndarrays, which will appropriately and explicitly
+        serve as constants elsewhere in a computational graph.
+
+        >>> x = mg.tensor([1., 2.])
+        >>> np.less_equal(x, 1)
+        array([ True, False])
+        """
+        out = kwargs.pop("out", (None,))
+        if len(out) > 1:  # pragma: no cover
+            raise ValueError(
+                "mygrad does not support in-place operations with more that one target"
+            )
+        (out,) = out
+
+        out: Optional[Union[np.ndarray, "Tensor"]]
+
+        try:
+            # differentiable ufunc implemented by mygrad
+            return getattr(_REGISTERED_UFUNC[ufunc], method)(*inputs, **kwargs, out=out)
+        except KeyError:
+            pass
+
+        # non-differentiable ufuncs get called on numpy arrays stored by tensors
+        if ufunc in _REGISTERED_BOOL_ONLY_UFUNC:
+            caster = asarray
+        elif ufunc in _REGISTERED_CONST_ONLY_UFUNC:
+            # the presence of non-constant tensors will raise
+            caster = _as_constant_array
+        else:  # pragma: no cover
+            return NotImplemented
+
+        try:
+            if out is not None:
+                kwargs["out"] = caster(out)
+            # returns ndarray
+            return getattr(ufunc, method)(*(caster(t) for t in inputs), **kwargs)
+        except _ConstantOnly:
+            raise ValueError(
+                f"{repr(ufunc)} cannot involve non-constant mygrad tensors."
+            )
+
     def __array__(self, dtype: DTypeLike = None) -> np.ndarray:
         return np.array(self.data, dtype=dtype, copy=False)
 
@@ -787,11 +930,25 @@ def _op(
         -------
         mygrad.Tensor
             The tensor-result of the operation's forward-pass."""
-        if out is not None and isinstance(out, Tensor):
-            out._in_place_op(
-                Op, *input_vars, op_args=op_args, op_kwargs=op_kwargs, constant=constant
-            )
-            return out
+        if out is not None:
+            if isinstance(out, tuple):
+                if len(out) > 1:  # pragma: no cover
+                    raise ValueError(
+                        "mygrad does not support in-place operations with more that one target"
+                    )
+                (out,) = out
+
+            if isinstance(out, Tensor):
+                out._in_place_op(
+                    Op,
+                    *input_vars,
+                    op_args=op_args,
+                    op_kwargs=op_kwargs,
+                    constant=constant,
+                )
+                return out
+
+        out: Optional[np.ndarray]
 
         _uniques_bases_then_arrs = ()
 
@@ -1700,21 +1857,11 @@ def __truediv__(self, other: ArrayLike) -> "Tensor":
     def __rtruediv__(self, other: ArrayLike) -> "Tensor":
         return self._op(Divide, other, self)
 
-    def __floordiv__(self, other: ArrayLike) -> "Tensor":
-        if not self.constant:
-            raise ValueError(
-                "Floor division cannot involve non-constant mygrad tensors."
-            )
-        if isinstance(other, Tensor):
-            other = other.data
-        return type(self)(self.data.__floordiv__(other), constant=True)
+    def __floordiv__(self, other: ArrayLike) -> np.ndarray:
+        return np.floor_divide(self, other)
 
-    def __rfloordiv__(self, other: ArrayLike) -> "Tensor":
-        if not self.constant:
-            raise ValueError(
-                "Floor division cannot involve non-constant mygrad tensors."
-            )
-        return type(self)(self.data.__rfloordiv__(other), constant=True)
+    def __rfloordiv__(self, other: ArrayLike) -> np.ndarray:
+        return np.floor_divide(other, self)
 
     def __itruediv__(self, other: ArrayLike) -> "Tensor":
         self._in_place_op(Divide, self, other)

src/mygrad/ufuncs/_ufunc_creators.py

@@ -16,8 +16,8 @@
 
 import numpy as np
 
-from mygrad import Tensor
 from mygrad.operation_base import BinaryUfunc, Operation, Ufunc, UnaryUfunc, _NoValue
+from mygrad.tensor_base import _REGISTERED_UFUNC, Tensor
 from mygrad.typing import ArrayLike, DTypeLikeReals, Index, Mask, Real
 
 __all__ = ["ufunc_creator"]
@@ -316,7 +316,7 @@ def _create_ufunc(
     outer_op=None,
     reduce_op=None,
     reduceat_op=None,
-):
+) -> Type[ufunc]:
     def at(
         a: ArrayLike,
         indices: Union[ArrayLike, Index, Tuple[ArrayLike, Index]],
@@ -330,7 +330,7 @@ def at(
     def accumulate(
         array: ArrayLike,
         axis: int = 0,
-        dtype: DTypeLikeReals = None,
+        dtype: Optional[DTypeLikeReals] = None,
         out: Optional[Union[Tensor, np.ndarray]] = None,
         *,
         constant: Optional[bool] = None,
@@ -342,16 +342,16 @@ def outer(
         a: ArrayLike,
         b: ArrayLike,
         *,
-        dtype: DTypeLikeReals,
-        out: Optional[Union[Tensor, np.ndarray]],
+        dtype: Optional[DTypeLikeReals] = None,
+        out: Optional[Union[Tensor, np.ndarray]] = None,
     ) -> Tensor:  # pragma: no cover
         """Not Implemented"""
         raise NotImplementedError()
 
     def reduce(
         a: ArrayLike,
         axis: Optional[Union[int, Tuple[int, ...]]] = 0,
-        dtype: DTypeLikeReals = None,
+        dtype: Optional[DTypeLikeReals] = None,
         out: Optional[Union[Tensor, np.ndarray]] = None,
         keepdims: bool = False,
         initial: Real = _NoValue,
@@ -364,7 +364,7 @@ def reduceat(
         a: ArrayLike,
         indices: ArrayLike,
         axis: Optional[Union[int, Tuple[int, ...]]] = 0,
-        dtype: DTypeLikeReals = None,
+        dtype: Optional[DTypeLikeReals] = None,
         out: Optional[Union[Tensor, np.ndarray]] = None,
     ) -> Tensor:  # pragma: no cover
         """Not Implemented"""
@@ -378,7 +378,7 @@ def reduceat(
         )
     else:  # pragma: no cover
         raise NotImplementedError(
-            "MyGrad Internal: `mygrad._utils.op_creator` only supports unary and binary ufuncs currently"
+            "MyGrad Internal: `mygrad._utils.op_creator` only supports unary and binary ufuncs presently"
         )
 
     # filter out non-real dtypes
@@ -487,7 +487,7 @@ def __init__(
         self.reduceat_op = reduceat_op
 
     def __call__(self, decorated_func: T) -> T:
-        return _create_ufunc(
+        out_ufunc = _create_ufunc(
             self.op,
             decorated_func=decorated_func,
             at_op=self.at_op,
@@ -496,3 +496,5 @@ def __call__(self, decorated_func: T) -> T:
             reduce_op=self.reduce_op,
             reduceat_op=self.reduceat_op,
         )
+        _REGISTERED_UFUNC[getattr(np, out_ufunc.__name__)] = out_ufunc
+        return out_ufunc

tests/tensor_base/test_operator_override.py

@@ -104,14 +104,16 @@ def test_arithmetic_operators_between_array_and_tensor_cast_to_tensor(
     "f1, f2",
     [
         (constant_tensor, lambda x: x),
-        (lambda x: x, constant_tensor),
+        (
+            lambda x: x.tolist(),
+            constant_tensor,
+        ),  # `list/tensor` ensures __rfloordiv__ gets called
         (constant_tensor, constant_tensor),
     ],
 )
 def test_floor_div(arr1, arr2, f1, f2):
     desired = arr1 // arr2
     actual = f1(arr1) // f2(arr2)
-    assert actual.constant is True
     assert actual.dtype == desired.dtype
     assert_array_equal(desired, actual)
 

tests/ufuncs/test_fwd_prop_and_backprop.py

@@ -216,15 +216,17 @@ def test_ufunc_fwd(
 @pytest.mark.parametrize(
     "ufunc", [u for u in ufuncs if u not in DOES_NOT_SUPPORT_COMPLEX_DOMAIN]
 )
-@given(data=st.data())
+@given(data=st.data(), use_numpy_overload=st.booleans())
 def test_ufunc_bkwd(
     data: st.DataObject,
     ufunc: Union[MyGradUnaryUfunc, MyGradBinaryUfunc],
+    use_numpy_overload: bool,
 ):
     """
     Checks:
     - backprop matches numerical gradient
     - backprop doesn't mutate grad
+    - that calling op through numpy overload works identically
     """
     args = data.draw(
         populates_ufunc(
@@ -237,6 +239,8 @@ def test_ufunc_bkwd(
     )
     args.make_array_based_args_read_only()  # guards against mutation
 
+    if use_numpy_overload:
+        ufunc = getattr(np, ufunc.__name__)
     mygrad_out = ufunc(*args.args, **args.kwargs)
 
     # Draw upstream gradient to be backpropped