Merge pull request #365 from rsokl/add-clip-method

add clip as method and add annotations

src/mygrad/__init__.py

@@ -35,3 +35,6 @@
 
 __version__ = get_versions()["version"]
 del get_versions
+
+
+setattr(Tensor, "clip", clip)

src/mygrad/math/misc/funcs.py

@@ -404,7 +404,9 @@ def minimum(
 
 
 @implements_numpy_override
-def clip(a, a_min, a_max, *, constant=None):
+def clip(
+    a: ArrayLike, a_min: ArrayLike, a_max: ArrayLike, *, constant: Optional[bool] = None
+) -> Tensor:
     """Clip (limit) the values in an array.
 
     Given an interval, values outside the interval are clipped to

src/mygrad/tensor_base.py

@@ -3071,3 +3071,56 @@ def any(
 
         """
         return np.any(self.data, axis=axis, out=out, keepdims=keepdims)
+
+    def clip(
+        self, a_min: ArrayLike, a_max: ArrayLike, *, constant: Optional[bool] = None
+    ) -> "Tensor":  # pragma: no cover
+        """Clip (limit) the values in an array.
+
+        Given an interval, values outside the interval are clipped to
+        the interval edges.  For example, if an interval of ``[0, 1]``
+        is specified, values smaller than 0 become 0, and values larger
+        than 1 become 1.
+
+        Equivalent to `mg.minimum(a_max, mg.maximum(a, a_min))``.
+
+        No check is performed to ensure ``a_min < a_max``.
+
+        This docstring was adapted from that of `numpy.clip`
+
+        Parameters
+        ----------
+        a_min : Optional[float, ArrayLike]
+            Minimum value. If `None`, clipping is not performed on lower
+            interval edge. Not more than one of `a_min` and `a_max` may be
+            `None`.
+
+        a_max : Optional[float, ArrayLike]
+            Maximum value. If `None`, clipping is not performed on upper
+            interval edge. Not more than one of `a_min` and `a_max` may be
+            `None`. If `a_min` or `a_max` are ArrayLike, then the three
+            arrays will be broadcasted to match their shapes.
+
+        constant : bool, optional(default=False)
+            If ``True``, the returned tensor is a constant (it
+            does not backpropagate a gradient)
+
+        Returns
+        -------
+        Tensor
+            A tensor with the elements of `a`, but where values
+            < `a_min` are replaced with `a_min`, and those > `a_max`
+            with `a_max`.
+
+        Examples
+        --------
+        >>> import mygrad as mg
+        >>> a = mg.arange(10)
+        >>> a
+        Tensor([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
+        >>> a.clip(1, 8)
+        Tensor([1, 1, 2, 3, 4, 5, 6, 7, 8, 8])
+        >>> a.clip([3, 4, 1, 1, 1, 4, 4, 4, 4, 4], 8)
+        Tensor([3, 4, 2, 3, 4, 5, 6, 7, 8, 8])"""
+        # set in added in mygrad.__init__
+        ...

tests/math/test_misc.py

@@ -8,6 +8,7 @@
 from hypothesis import given, settings
 from numpy.testing import assert_allclose
 
+import mygrad as mg
 from mygrad import clip, maximum, minimum
 from mygrad.tensor_base import Tensor
 from tests.wrappers.uber import backprop_test_factory, fwdprop_test_factory
@@ -204,3 +205,16 @@ def test_clip_input_validation(a, a_min, a_max):
     mygrad_out = clip(a, a_min, a_max)
 
     np.testing.assert_array_equal(numpy_out, mygrad_out.data)
+
+
+def test_clip_method_fwd():
+    a = mg.arange(10.0)
+    assert_allclose(
+        a.clip([3, 4, 1, 1, 1, 4, 4, 4, 4, 4], 8), [3, 4, 2, 3, 4, 5, 6, 7, 8, 8]
+    )
+
+
+def test_clip_method_bkwd():
+    x = mg.tensor([1.0, 5.0, 10.0])
+    x.clip(2, 7).backward()
+    assert_allclose(x.grad, [0.0, 1.0, 0.0])