add unbind python api (#24141)

* add unbind pyhon api. test=develp

python/paddle/fluid/layers/nn.py

@@ -186,6 +186,7 @@ __all__ = [
     'hard_swish',
     'gather_tree',
     'uniform_random',
+    'unbind',
 ]
 
 
@@ -14344,3 +14345,57 @@ def uniform_random(shape, dtype='float32', min=-1.0, max=1.0, seed=0):
         outputs={"Out": out})
 
     return helper.append_activation(out)
+
+
+def unbind(input, axis=0):
+    """
+    Removes a tensor dimension, then split the input tensor into multiple sub-Tensors.
+    Args:
+        input (Variable): The input variable which is an N-D Tensor, data type being float32, float64, int32 or int64.
+       
+        axis (int32|int64, optional): A scalar with type ``int32|int64`` shape [1]. The dimension along which to unbind. If :math:`axis < 0`, the
+            dimension to unbind along is :math:`rank(input) + axis`. Default is 0.
+    Returns:
+        list(Variable): The list of segmented Tensor variables.
+
+    Example:
+        .. code-block:: python
+            import paddle
+            # input is a variable which shape is [3, 4, 5]
+            input = paddle.fluid.data(
+                 name="input", shape=[3, 4, 5], dtype="float32")
+            [x0, x1, x2] = paddle.tensor.unbind(input, axis=0)
+            # x0.shape [4, 5]
+            # x1.shape [4, 5]
+            # x2.shape [4, 5]
+            [x0, x1, x2, x3] = paddle.tensor.unbind(input, axis=1)
+            # x0.shape [3, 5]
+            # x1.shape [3, 5]
+            # x2.shape [3, 5]
+            # x3.shape [3, 5]
+
+    """
+    helper = LayerHelper("unbind", **locals())
+    check_type(input, 'input', (Variable), 'unbind')
+    dtype = helper.input_dtype()
+    check_dtype(dtype, 'unbind', ['float32', 'float64', 'int32', 'int64'],
+                'unbind')
+    if not isinstance(axis, (int)):
+        raise TypeError("The type of 'axis'  must be int, but received %s." %
+                        (type(axis)))
+    if isinstance(axis, np.generic):
+        axis = np.asscalar(axis)
+    input_shape = input.shape
+    axis_ = axis if axis >= 0 else len(input_shape) + axis
+    num = input_shape[axis_]
+    outs = [
+        helper.create_variable_for_type_inference(dtype=helper.input_dtype())
+        for i in range(num)
+    ]
+
+    helper.append_op(
+        type="unbind",
+        inputs={"X": input},
+        outputs={"Out": outs},
+        attrs={"axis": axis})
+    return outs

python/paddle/fluid/tests/unittests/test_unbind_op.py

@@ -18,9 +18,46 @@ import unittest
 import numpy as np
 from op_test import OpTest
 import paddle.fluid as fluid
+import paddle.tensor as tensor
 from paddle.fluid import compiler, Program, program_guard, core
 
 
+class TestUnbind(unittest.TestCase):
+    def test_unbind(self):
+
+        x_1 = fluid.data(shape=[2, 3], dtype='float32', name='x_1')
+        [out_0, out_1] = tensor.unbind(input=x_1, axis=0)
+        input_1 = np.random.random([2, 3]).astype("float32")
+        axis = fluid.data(shape=[1], dtype='int32', name='axis')
+        exe = fluid.Executor(place=fluid.CPUPlace())
+
+        [res_1, res_2] = exe.run(fluid.default_main_program(),
+                                 feed={"x_1": input_1,
+                                       "axis": 0},
+                                 fetch_list=[out_0, out_1])
+
+        assert np.array_equal(res_1, input_1[0, 0:100])
+        assert np.array_equal(res_2, input_1[1, 0:100])
+
+
+class TestLayersUnbind(unittest.TestCase):
+    def test_layers_unbind(self):
+
+        x_1 = fluid.data(shape=[2, 3], dtype='float32', name='x_1')
+        [out_0, out_1] = fluid.layers.unbind(input=x_1, axis=0)
+        input_1 = np.random.random([2, 3]).astype("float32")
+        axis = fluid.data(shape=[1], dtype='int32', name='axis')
+        exe = fluid.Executor(place=fluid.CPUPlace())
+
+        [res_1, res_2] = exe.run(fluid.default_main_program(),
+                                 feed={"x_1": input_1,
+                                       "axis": 0},
+                                 fetch_list=[out_0, out_1])
+
+        assert np.array_equal(res_1, input_1[0, 0:100])
+        assert np.array_equal(res_2, input_1[1, 0:100])
+
+
 class TestUnbindOp(OpTest):
     def initParameters(self):
         pass
@@ -37,7 +74,6 @@ class TestUnbindOp(OpTest):
         self.axis = 0
         self.num = 3
         self.initParameters()
-        #x = np.random.random((3, 2, 2)).astype(self.dtype)
         x = np.arange(12).reshape(3, 2, 2).astype(self.dtype)
         self.out = np.split(x, self.num, self.axis)
         self.outReshape()
@@ -118,5 +154,16 @@ class TestUnbindOp4(TestUnbindOp):
         self.out[1] = self.out[1].reshape((3, 2))
 
 
+class TestUnbindAxisError(unittest.TestCase):
+    def test_errors(self):
+        with program_guard(Program(), Program()):
+            x = fluid.data(shape=[2, 3], dtype='float32', name='x')
+
+            def test_table_Variable():
+                tensor.unbind(input=x, axis=2.0)
+
+            self.assertRaises(TypeError, test_table_Variable)
+
+
 if __name__ == '__main__':
     unittest.main()

python/paddle/tensor/__init__.py

@@ -151,6 +151,7 @@ from .math import erf  #DEFINE_ALIAS
 from .math import addcmul  #DEFINE_ALIAS
 from .math import addmm  #DEFINE_ALIAS
 from .math import clamp  #DEFINE_ALIAS
+from .manipulation import unbind  #DEFINE_ALIAS
 from .math import trace  #DEFINE_ALIAS
 from .math import kron  #DEFINE_ALIAS
 # from .random import gaussin        #DEFINE_ALIAS

python/paddle/tensor/manipulation.py

@@ -20,6 +20,7 @@ from ..fluid.framework import Variable, OpProtoHolder, in_dygraph_mode, convert_
 from ..fluid.data_feeder import convert_dtype, check_variable_and_dtype, check_type, check_dtype
 from ..fluid.layers.tensor import fill_constant
 from ..fluid.layers import utils
+import numpy as np
 # TODO: define functions to manipulate a tensor  
 from ..fluid.layers import cast  #DEFINE_ALIAS
 from ..fluid.layers import concat  #DEFINE_ALIAS
@@ -60,7 +61,7 @@ __all__ = [
     'unsqueeze',
     'unstack',
     'flip',
-    #       'unbind',
+    'unbind',
     'roll'
 ]
 
@@ -657,3 +658,57 @@ def gather(input, index, overwrite=True):
         outputs={"Out": out},
         attrs={'overwrite': overwrite})
     return out
+
+
+def unbind(input, axis=0):
+    """
+    Removes a tensor dimension, then split the input tensor into multiple sub-Tensors.
+    Args:
+        input (Variable): The input variable which is an N-D Tensor, data type being float32, float64, int32 or int64.
+       
+        axis (int32|int64, optional): A scalar with type ``int32|int64`` shape [1]. The dimension along which to unbind. If :math:`axis < 0`, the
+            dimension to unbind along is :math:`rank(input) + axis`. Default is 0.
+    Returns:
+        list(Variable): The list of segmented Tensor variables.
+
+    Example:
+        .. code-block:: python
+            import paddle
+            # input is a variable which shape is [3, 4, 5]
+            input = paddle.fluid.data(
+                 name="input", shape=[3, 4, 5], dtype="float32")
+            [x0, x1, x2] = paddle.tensor.unbind(input, axis=0)
+            # x0.shape [4, 5]
+            # x1.shape [4, 5]
+            # x2.shape [4, 5]
+            [x0, x1, x2, x3] = paddle.tensor.unbind(input, axis=1)
+            # x0.shape [3, 5]
+            # x1.shape [3, 5]
+            # x2.shape [3, 5]
+            # x3.shape [3, 5]
+
+    """
+    helper = LayerHelper("unbind", **locals())
+    check_type(input, 'input', (Variable), 'unbind')
+    dtype = helper.input_dtype()
+    check_dtype(dtype, 'unbind', ['float32', 'float64', 'int32', 'int64'],
+                'unbind')
+    if not isinstance(axis, (int)):
+        raise TypeError("The type of 'axis'  must be int, but received %s." %
+                        (type(axis)))
+    if isinstance(axis, np.generic):
+        axis = np.asscalar(axis)
+    input_shape = input.shape
+    axis_ = axis if axis >= 0 else len(input_shape) + axis
+    num = input_shape[axis_]
+    outs = [
+        helper.create_variable_for_type_inference(dtype=helper.input_dtype())
+        for i in range(num)
+    ]
+
+    helper.append_op(
+        type="unbind",
+        inputs={"X": input},
+        outputs={"Out": outs},
+        attrs={"axis": axis})
+    return outs