Add convenience fluent method for squeeze op (apache#9734)

* Add convenience fluent entry for squeeze op * Add unit tests
rahul003 · Feb 9, 2018 · 54c7989 · 54c7989
1 parent 333f2f2
commit 54c7989
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 0 deletions.
diff --git a/python/mxnet/ndarray/ndarray.py b/python/mxnet/ndarray/ndarray.py
@@ -1545,6 +1545,14 @@ def log_softmax(self, *args, **kwargs):
         """
         return op.log_softmax(self, *args, **kwargs)
 
+    def squeeze(self, *args, **kwargs):
+        """Convenience fluent method for :py:func:`squeeze`.
+
+        The arguments are the same as for :py:func:`squeeze`, with
+        this array as data.
+        """
+        return op.squeeze(self, *args, **kwargs)
+
     # pylint: disable= undefined-variable
     def broadcast_to(self, shape):
         """Broadcasts the input array to a new shape.

diff --git a/python/mxnet/symbol/symbol.py b/python/mxnet/symbol/symbol.py
@@ -2356,6 +2356,14 @@ def log_softmax(self, *args, **kwargs):
         """
         return op.log_softmax(self, *args, **kwargs)
 
+    def squeeze(self, *args, **kwargs):
+        """Convenience fluent method for :py:func:`squeeze`.
+
+        The arguments are the same as for :py:func:`squeeze`, with
+        this array as data.
+        """
+        return op.squeeze(self, *args, **kwargs)
+
     def wait_to_read(self):
         raise NotImplementedForSymbol(self.wait_to_read, None)
 

diff --git a/tests/python/unittest/test_ndarray.py b/tests/python/unittest/test_ndarray.py
@@ -790,6 +790,7 @@ def check_fluent_regular(func, kwargs, shape=(5, 17, 1), equal_nan=False):
 
     check_fluent_regular('reshape', {'shape': (17, 1, 5)})
     check_fluent_regular('broadcast_to', {'shape': (5, 17, 47)})
+    check_fluent_regular('squeeze', {'axis': (1, 3)}, shape=(2, 1, 3, 1, 4))
 
 @raises(ValueError)
 def test_bool_ambiguous():

diff --git a/tests/python/unittest/test_symbol.py b/tests/python/unittest/test_symbol.py
@@ -214,6 +214,7 @@ def check_fluent_regular(func, kwargs, shape=(5, 17, 1), equal_nan=False):
 
     check_fluent_regular('reshape', {'shape': (17, 1, 5)})
     check_fluent_regular('broadcast_to', {'shape': (5, 17, 47)})
+    check_fluent_regular('squeeze', {'axis': (1, 3)}, shape=(2, 1, 3, 1, 4))
 
 def check_symbol_consistency(sym1, sym2, ctx, skip_grad=False, equal_nan=False):
     assert sym1.list_arguments() == sym2.list_arguments()