Add convenience fluent method for squeeze op

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.

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)
 

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():

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()