test_split_op.py 25.7 KB
Newer Older
1
#   Copyright (c) 2018 PaddlePaddle Authors. All Rights Reserved.
D
dzhwinter 已提交
2
#
D
dzhwinter 已提交
3 4 5
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
D
dzhwinter 已提交
6
#
D
dzhwinter 已提交
7
#     http://www.apache.org/licenses/LICENSE-2.0
D
dzhwinter 已提交
8
#
D
dzhwinter 已提交
9 10 11 12 13 14
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

Y
Yancey 已提交
15
import unittest
16

Y
Yancey 已提交
17
import numpy as np
18
from eager_op_test import OpTest, convert_float_to_uint16
19 20

import paddle
21
from paddle import fluid
22
from paddle.fluid import Program, core, program_guard
Y
Yancey 已提交
23 24 25 26


class TestSplitOp(OpTest):
    def setUp(self):
27
        self.python_api = paddle.split
28
        self.public_python_api = paddle.split
29
        self.python_out_sig = ['out0', 'out1', 'out2']
T
fix ut  
typhoonzero 已提交
30
        self._set_op_type()
31
        self.prim_op_type = "prim"
32
        self.dtype = self.get_dtype()
Y
Yancey1989 已提交
33
        axis = 1
34 35 36 37
        if self.dtype == np.uint16:
            x = np.random.random((4, 5, 6)).astype(np.float32)
            out = np.split(x, [2, 3], axis)
            self.inputs = {'X': convert_float_to_uint16(x)}
38 39 40 41 42 43
            self.outputs = {
                'Out': [
                    ('out%d' % i, convert_float_to_uint16(out[i]))
                    for i in range(len(out))
                ]
            }
44 45 46 47
        else:
            x = np.random.random((4, 5, 6)).astype(self.dtype)
            out = np.split(x, [2, 3], axis)
            self.inputs = {'X': x}
48 49 50
            self.outputs = {
                'Out': [('out%d' % i, out[i]) for i in range(len(out))]
            }
Y
Yancey1989 已提交
51
        self.attrs = {'axis': axis, 'sections': [2, 1, 2]}
Y
Yancey 已提交
52

53
    def get_dtype(self):
54
        return "float64"
55

T
typhoonzero 已提交
56 57 58
    def _set_op_type(self):
        self.op_type = "split"

Y
Yancey 已提交
59 60 61
    def test_check_output(self):
        self.check_output()

Y
Yancey1989 已提交
62
    def test_check_grad(self):
63
        self.check_grad(['X'], ['out0', 'out1', 'out2'], check_prim=True)
Y
Yancey 已提交
64 65


66
# test with attr(num)
67
class TestSplitWithNumOp(OpTest):
68
    def setUp(self):
69
        self.python_api = paddle.split
70
        self.public_python_api = paddle.split
71
        self.python_out_sig = ['out0', 'out1', 'out2']
72
        self._set_op_type()
73
        self.prim_op_type = "prim"
74 75 76 77 78
        self.dtype = self.get_dtype()
        self.init_data()
        self.attrs = {
            'axis': self.axis,
            'sections': self.sections,
79
            'num': self.num,
80
        }
81 82 83 84 85 86 87 88 89 90 91 92 93 94 95
        if self.dtype == np.uint16:
            self.inputs = {'X': convert_float_to_uint16(self.x)}
            out = np.split(self.x, self.indices_or_sections, self.axis)
            self.outputs = {
                'Out': [
                    ('out%d' % i, convert_float_to_uint16(out[i]))
                    for i in range(len(out))
                ]
            }
        else:
            self.inputs = {'X': self.x}
            out = np.split(self.x, self.indices_or_sections, self.axis)
            self.outputs = {
                'Out': [('out%d' % i, out[i]) for i in range(len(out))]
            }
96 97

    def init_data(self):
98 99 100 101
        if self.dtype == np.uint16:
            self.x = np.random.random((4, 5, 6)).astype(np.float32)
        else:
            self.x = np.random.random((4, 5, 6)).astype(self.dtype)
102 103 104 105 106 107
        self.axis = 2
        self.sections = []
        self.num = 3
        self.indices_or_sections = 3

    def get_dtype(self):
108
        return "float64"
109 110 111 112 113 114 115 116

    def _set_op_type(self):
        self.op_type = "split"

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
117
        self.check_grad(['X'], ['out0', 'out1', 'out2'], check_prim=True)
118 119 120 121 122


# attr(axis) is Tensor
class TestSplitOp_AxisTensor(OpTest):
    def setUp(self):
123 124
        self.python_api = paddle.split
        self.python_out_sig = ['out0', 'out1', 'out2']
125 126 127 128 129
        self._set_op_type()
        self.dtype = self.get_dtype()
        self.init_data()
        self.inputs = {
            'X': self.x,
130
            'AxisTensor': np.array([self.axis]).astype("int32"),
131 132 133 134
        }
        self.attrs = {'sections': self.sections, 'num': self.num}

        out = np.split(self.x, self.indices_or_sections, self.axis)
135
        self.outputs = {'Out': [('out%d' % i, out[i]) for i in range(len(out))]}
136 137 138 139 140 141 142 143 144

    def init_data(self):
        self.x = np.random.random((4, 5, 6)).astype(self.dtype)
        self.axis = 2
        self.sections = []
        self.num = 3
        self.indices_or_sections = 3

    def get_dtype(self):
145
        return "float64"
146 147 148 149 150 151 152 153 154 155 156 157 158 159

    def _set_op_type(self):
        self.op_type = "split"

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], ['out0', 'out1', 'out2'])


# attr(sections) is list containing Tensor
class TestSplitOp_SectionsTensor(OpTest):
    def setUp(self):
160 161
        self.python_api = paddle.split
        self.python_out_sig = ['out0', 'out1', 'out2']
162 163 164 165 166 167 168
        self._set_op_type()
        self.dtype = self.get_dtype()
        self.init_data()
        self.inputs = {'X': self.x}

        sections_tensor = []
        for index, ele in enumerate(self.sections):
169
            sections_tensor.append(
170
                ("x" + str(index), np.ones(1).astype('int32') * ele)
171
            )
172 173 174 175 176 177

        self.inputs['SectionsTensorList'] = sections_tensor

        self.attrs = {
            'axis': self.axis,
            'sections': self.sections_infer,
178
            'num': self.num,
179 180 181
        }

        out = np.split(self.x, self.indices_or_sections, self.axis)
182
        self.outputs = {'Out': [('out%d' % i, out[i]) for i in range(len(out))]}
183 184 185 186 187 188 189 190 191 192

    def init_data(self):
        self.x = np.random.random((4, 5, 6)).astype(self.dtype)
        self.axis = 1
        self.sections = [2, 1, 2]
        self.sections_infer = [-1, -1, -1]
        self.num = 0
        self.indices_or_sections = [2, 3]

    def get_dtype(self):
193
        return "float64"
194 195 196 197 198 199 200 201 202 203 204 205 206

    def _set_op_type(self):
        self.op_type = "split"

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], ['out0', 'out1', 'out2'])


class TestSplitOp_unk_section(OpTest):
    def setUp(self):
207
        self.python_api = paddle.split
208
        self.public_python_api = paddle.split
209
        self.python_out_sig = ['out0', 'out1', 'out2']
210
        self._set_op_type()
211
        self.prim_op_type = "prim"
212 213 214 215 216 217
        self.dtype = self.get_dtype()
        self.init_data()
        self.inputs = {'X': self.x}
        self.attrs = {
            'axis': self.axis,
            'sections': self.sections,
218
            'num': self.num,
219 220 221
        }

        out = np.split(self.x, self.indices_or_sections, self.axis)
222
        self.outputs = {'Out': [('out%d' % i, out[i]) for i in range(len(out))]}
223 224 225 226 227 228 229 230 231

    def init_data(self):
        self.x = np.random.random((4, 5, 6)).astype(self.dtype)
        self.axis = 2
        self.sections = [2, 1, -1]
        self.num = 0
        self.indices_or_sections = [2, 3]

    def get_dtype(self):
232
        return "float64"
233 234 235 236 237 238 239 240

    def _set_op_type(self):
        self.op_type = "split"

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
241
        self.check_grad(['X'], ['out0', 'out1', 'out2'], check_prim=True)
242 243


T
typhoonzero 已提交
244 245 246 247 248
class TestSplitByrefOp(OpTest):
    def _set_op_type(self):
        self.op_type = "split_byref"


249
# ----------------Split Fp16----------------
250 251 252


def create_test_fp16(parent):
253 254 255
    @unittest.skipIf(
        not core.is_compiled_with_cuda(), "core is not compiled with CUDA"
    )
256
    class TestSplitFP16Op(parent):
257 258 259
        def get_dtype(self):
            return np.float16

260 261 262
    cls_name = "{}_{}".format(parent.__name__, "FP16Op")
    TestSplitFP16Op.__name__ = cls_name
    globals()[cls_name] = TestSplitFP16Op
263 264 265


create_test_fp16(TestSplitOp)
266
create_test_fp16(TestSplitWithNumOp)
267

268
# ----------------Split Bf16----------------
269 270 271


def create_test_bf16(parent):
272
    @unittest.skipIf(
273 274 275
        not core.is_compiled_with_cuda()
        or not core.is_bfloat16_supported(core.CUDAPlace(0)),
        "core is not compiled with CUDA or not support bfloat16",
276
    )
277
    class TestSplitBF16Op(parent):
278 279 280 281 282 283 284 285
        def get_dtype(self):
            return np.uint16

        def test_check_output(self):
            place = core.CUDAPlace(0)
            self.check_output_with_place(place)

        def test_check_grad(self):
286
            place = core.CUDAPlace(0)
287
            self.check_grad_with_place(place, ['X'], 'out2', check_prim=True)
288

289 290 291
    cls_name = "{}_{}".format(parent.__name__, "BF16Op")
    TestSplitBF16Op.__name__ = cls_name
    globals()[cls_name] = TestSplitBF16Op
292 293 294


create_test_bf16(TestSplitOp)
295
create_test_bf16(TestSplitWithNumOp)
296

297

298
class TestSplitAPI(unittest.TestCase):
299 300
    def test_api(self):
        input_1 = np.random.random([4, 5, 6]).astype("int32")
301 302 303
        positive_1_int32 = paddle.tensor.fill_constant([1], "int32", 1)
        positive_1_int64 = paddle.tensor.fill_constant([1], "int64", 1)
        positive_2_int64 = paddle.tensor.fill_constant([1], "int64", 2)
304 305
        x_1 = paddle.static.data(shape=[4, 5, 6], dtype='int32', name='x_1')
        x_2 = paddle.static.data(shape=[4, 5, None], dtype='int32', name='x_2')
306

307 308
        out_0, out_1, out_2 = paddle.split(
            x=x_1,
309
            num_or_sections=[positive_2_int64, positive_1_int32, -1],
310
            axis=positive_1_int64,
311
        )
312

313 314
        out_3, out_4, out_5 = paddle.split(
            x=x_1, num_or_sections=[2, 1, 2], axis=positive_1_int32
315
        )
316
        paddle.split(x=x_2, num_or_sections=2, axis=2)
317 318

        exe = fluid.Executor(place=fluid.CPUPlace())
319 320 321 322 323
        [res_0, res_1, res_2, res_3, res_4, res_5] = exe.run(
            fluid.default_main_program(),
            feed={"x_1": input_1, "x_2": input_1},
            fetch_list=[out_0, out_1, out_2, out_3, out_4, out_5],
        )
324 325 326 327 328 329 330 331 332 333

        out = np.split(input_1, [2, 3], 1)
        assert np.array_equal(res_0, out[0])
        assert np.array_equal(res_1, out[1])
        assert np.array_equal(res_2, out[2])
        assert np.array_equal(res_3, out[0])
        assert np.array_equal(res_4, out[1])
        assert np.array_equal(res_5, out[2])


334
class TestSplitOpError(unittest.TestCase):
335 336 337 338
    def test_errors(self):
        with program_guard(Program(), Program()):
            # The type of axis in split_op should be int or Variable.
            def test_axis_type():
G
GGBond8488 已提交
339 340 341
                x6 = paddle.static.data(
                    shape=[-1, 4], dtype='float16', name='x3'
                )
342
                paddle.split(x=x6, num_or_sections=2, axis=3.2)
343 344 345

            self.assertRaises(TypeError, test_axis_type)

346 347
            # The type of axis in split_op should be int or Variable.
            def test_axis_variable_type():
G
GGBond8488 已提交
348 349 350 351 352 353
                x9 = paddle.static.data(
                    shape=[-1, 4], dtype='float16', name='x9'
                )
                x10 = paddle.static.data(
                    shape=[-1, 1], dtype='float16', name='x10'
                )
354
                paddle.split(x=x9, num_or_sections=2, axis=x10)
355 356 357

            self.assertRaises(TypeError, test_axis_variable_type)

358 359
            # The type of num_or_sections in split_op should be int, tuple or list.
            def test_num_or_sections_type():
G
GGBond8488 已提交
360 361 362
                x6 = paddle.static.data(
                    shape=[-1, 4], dtype='float16', name='x4'
                )
363
                paddle.split(x=x6, num_or_sections=2.1, axis=3)
364 365 366

            self.assertRaises(TypeError, test_num_or_sections_type)

367
            def test_num_or_sections_type_tensor():
G
GGBond8488 已提交
368 369 370
                x7 = paddle.static.data(
                    shape=[-1, 4], dtype='float16', name='x5'
                )
371 372 373 374 375
                paddle.split(input=x7, num_or_sections=2.1, dim=3)

            self.assertRaises(TypeError, test_num_or_sections_type_tensor)

            def test_axis_type_tensor():
G
GGBond8488 已提交
376 377 378
                x8 = paddle.static.data(
                    shape=[-1, 4], dtype='float16', name='x6'
                )
379 380 381 382
                paddle.split(input=x8, num_or_sections=2, dim=3.2)

            self.assertRaises(TypeError, test_axis_type_tensor)

张春乔 已提交
383 384 385 386 387 388 389 390
        with paddle.fluid.dygraph.guard():

            def test_0_num_tensor():
                x = paddle.uniform([1, 1, 1], dtype='float32')
                paddle.split(x, num_or_sections=0)

            self.assertRaises(ValueError, test_0_num_tensor)

391 392 393 394

class API_TestSplit(unittest.TestCase):
    def test_out(self):
        with fluid.program_guard(fluid.Program(), fluid.Program()):
G
GGBond8488 已提交
395 396 397 398 399 400
            data1 = paddle.static.data(
                'data1', shape=[-1, 4, 6, 6], dtype='float64'
            )
            data1.desc.set_need_check_feed(False)
            data2 = paddle.static.data('data2', shape=[-1, 1], dtype='int32')
            data2.desc.set_need_check_feed(False)
401
            x0, x1, x2 = paddle.split(data1, num_or_sections=3, axis=data2)
402 403 404 405
            place = fluid.CPUPlace()
            exe = fluid.Executor(place)
            input1 = np.random.random([4, 6, 6]).astype('float64')
            input2 = np.array([2]).astype('int32')
406 407 408
            r0, r1, r2, = exe.run(
                feed={"data1": input1, "data2": input2}, fetch_list=[x0, x1, x2]
            )
409
            ex_x0, ex_x1, ex_x2 = np.split(input1, 3, axis=2)
410 411 412
            np.testing.assert_allclose(ex_x0, r0, rtol=1e-05)
            np.testing.assert_allclose(ex_x1, r1, rtol=1e-05)
            np.testing.assert_allclose(ex_x2, r2, rtol=1e-05)
413 414 415 416 417


class API_TestSplit2(unittest.TestCase):
    def test_out(self):
        with fluid.program_guard(fluid.Program(), fluid.Program()):
G
GGBond8488 已提交
418 419 420 421
            data1 = paddle.static.data(
                'data1', shape=[-1, 4, 6, 6], dtype='float64'
            )
            data1.desc.set_need_check_feed(False)
422
            x0, x1, x2 = paddle.split(data1, num_or_sections=3, axis=2)
423 424 425
            place = fluid.CPUPlace()
            exe = fluid.Executor(place)
            input1 = np.random.random([4, 6, 6]).astype('float64')
426 427 428 429 430
            (
                r0,
                r1,
                r2,
            ) = exe.run(feed={"data1": input1}, fetch_list=[x0, x1, x2])
431
            ex_x0, ex_x1, ex_x2 = np.split(input1, 3, axis=2)
432 433 434
            np.testing.assert_allclose(ex_x0, r0, rtol=1e-05)
            np.testing.assert_allclose(ex_x1, r1, rtol=1e-05)
            np.testing.assert_allclose(ex_x2, r2, rtol=1e-05)
435 436 437 438 439


class API_TestSplit3(unittest.TestCase):
    def test_out(self):
        with fluid.program_guard(fluid.Program(), fluid.Program()):
G
GGBond8488 已提交
440
            data = paddle.static.data('data', shape=[-1, 10], dtype='float64')
441
            x0, x1 = paddle.split(data, num_or_sections=(3, 7), axis=1)
442 443 444 445
            place = fluid.CPUPlace()
            exe = fluid.Executor(place)
            input1 = np.random.random([1, 10]).astype('float64')
            r0, r1 = exe.run(feed={"data": input1}, fetch_list=[x0, x1])
446
            ex_x0, ex_x1 = np.split(input1, (3,), axis=1)
447 448
            np.testing.assert_allclose(ex_x0, r0, rtol=1e-05)
            np.testing.assert_allclose(ex_x1, r1, rtol=1e-05)
449 450 451 452 453


class API_TestSplit4(unittest.TestCase):
    def test_out(self):
        with fluid.program_guard(fluid.Program(), fluid.Program()):
G
GGBond8488 已提交
454 455
            data = paddle.static.data('data', shape=[-1, 10], dtype='float64')
            index = paddle.static.data('index', shape=[1], dtype='int32')
456
            x0, x1 = paddle.split(data, num_or_sections=(3, index), axis=1)
457 458 459 460
            place = fluid.CPUPlace()
            exe = fluid.Executor(place)
            input1 = np.random.random([1, 10]).astype('float64')
            input2 = np.array([7]).astype('int32')
461 462 463 464
            r0, r1 = exe.run(
                feed={"data": input1, "index": input2}, fetch_list=[x0, x1]
            )
            ex_x0, ex_x1 = np.split(input1, (3,), axis=1)
465 466
            np.testing.assert_allclose(ex_x0, r0, rtol=1e-05)
            np.testing.assert_allclose(ex_x1, r1, rtol=1e-05)
467 468


C
Charles-hit 已提交
469 470
class API_TestSplit5(unittest.TestCase):
    def test_out(self):
471 472 473
        for use_cuda in (
            [False, True] if core.is_compiled_with_cuda() else [False]
        ):
C
Charles-hit 已提交
474 475 476 477 478 479 480 481 482 483 484 485 486 487 488
            place = paddle.CUDAPlace(0) if use_cuda else paddle.CPUPlace()
            with fluid.program_guard(fluid.Program(), fluid.Program()):
                input_1 = np.random.random([5, 4]).astype("int32")
                # input is a variable which shape is [5, 4]
                input = paddle.to_tensor(input_1)
                n = paddle.full([1], 5, dtype='int32')
                out = paddle.split(input, [n])
                exe = paddle.static.Executor(place=place)
                re = exe.run(fetch_list=[out])
                re = re[0]
                ex_out = np.split(input_1, [5])
                ex_out = ex_out[0]
                np.testing.assert_allclose(ex_out, re, rtol=1e-05)


489 490 491
class API_TestSplit6(unittest.TestCase):
    def test_out(self):
        with fluid.program_guard(fluid.Program(), fluid.Program()):
G
GGBond8488 已提交
492
            data = paddle.static.data('data', shape=[-1, 10], dtype='float64')
493 494 495 496 497
            x0, x1 = paddle.split(data, num_or_sections=[1, 1], axis=0)
            place = fluid.CPUPlace()
            exe = fluid.Executor(place)
            input1 = np.random.random([2, 10]).astype('float64')
            r0, r1 = exe.run(feed={"data": input1}, fetch_list=[x0, x1])
498
            ex_x0, ex_x1 = np.split(input1, (1,), axis=0)
499 500 501 502
            np.testing.assert_allclose(ex_x0, r0, rtol=1e-05)
            np.testing.assert_allclose(ex_x1, r1, rtol=1e-05)


C
Charles-hit 已提交
503 504 505 506 507 508
class API_TestDygraphFluidSplit(unittest.TestCase):
    def test_out1(self):
        with fluid.dygraph.guard():
            input_1 = np.random.random([4, 6, 6]).astype("int32")
            # input is a variable which shape is [4, 6, 6]
            input = paddle.to_tensor(input_1)
509
            x0, x1, x2 = paddle.split(input, num_or_sections=3, axis=1)
C
Charles-hit 已提交
510 511 512 513
            x0_out = x0.numpy()
            x1_out = x1.numpy()
            x2_out = x2.numpy()
            ex_x0, ex_x1, ex_x2 = np.split(input_1, 3, axis=1)
514 515 516
            # input is a variable which shape is [4, 6, 6]
            input = paddle.to_tensor(input_1)
            input.stop_gradient = False
517
            x0, x1, x2 = paddle.split(input, num_or_sections=3, axis=1)
518 519 520 521 522 523 524 525 526 527 528
            eager_x0_out = x0.numpy()
            eager_x1_out = x1.numpy()
            eager_x2_out = x2.numpy()
            loss = x0.sum()
            loss.backward()
            manul_grad = np.zeros_like(input_1)
            manul_grad[:, :2, :] = 1
            np.testing.assert_allclose(input.gradient(), manul_grad, rtol=1e-05)
            np.testing.assert_allclose(ex_x0, eager_x0_out, rtol=1e-05)
            np.testing.assert_allclose(ex_x1, eager_x1_out, rtol=1e-05)
            np.testing.assert_allclose(ex_x2, eager_x2_out, rtol=1e-05)
C
Charles-hit 已提交
529 530 531 532 533 534 535 536 537 538

        np.testing.assert_allclose(ex_x0, x0_out, rtol=1e-05)
        np.testing.assert_allclose(ex_x1, x1_out, rtol=1e-05)
        np.testing.assert_allclose(ex_x2, x2_out, rtol=1e-05)

    def test_out2(self):
        with fluid.dygraph.guard():
            input_1 = np.random.random([4, 6, 6]).astype("int32")
            # input is a variable which shape is [4, 6, 6]
            input = paddle.to_tensor(input_1)
539
            x0, x1, x2 = paddle.split(input, [2, 2, 2], axis=1)
C
Charles-hit 已提交
540 541 542 543
            x0_out = x0.numpy()
            x1_out = x1.numpy()
            x2_out = x2.numpy()
            ex_x0, ex_x1, ex_x2 = np.split(input_1, 3, axis=1)
544 545 546
            # input is a variable which shape is [4, 6, 6]
            input = paddle.to_tensor(input_1)
            input.stop_gradient = False
547
            x0, x1, x2 = paddle.split(input, [2, 2, 2], axis=1)
548 549 550 551 552 553 554 555 556 557 558
            eager_x0_out = x0.numpy()
            eager_x1_out = x1.numpy()
            eager_x2_out = x2.numpy()
            loss = x0.sum()
            loss.backward()
            manul_grad = np.zeros_like(input_1)
            manul_grad[:, :2, :] = 1
            np.testing.assert_allclose(input.gradient(), manul_grad, rtol=1e-05)
            np.testing.assert_allclose(ex_x0, eager_x0_out, rtol=1e-05)
            np.testing.assert_allclose(ex_x1, eager_x1_out, rtol=1e-05)
            np.testing.assert_allclose(ex_x2, eager_x2_out, rtol=1e-05)
C
Charles-hit 已提交
559 560 561 562 563 564

        np.testing.assert_allclose(ex_x0, x0_out, rtol=1e-05)
        np.testing.assert_allclose(ex_x1, x1_out, rtol=1e-05)
        np.testing.assert_allclose(ex_x2, x2_out, rtol=1e-05)


565
class API_TestDygraphSplit(unittest.TestCase):
566 567 568 569
    def test_out1(self):
        with fluid.dygraph.guard():
            input_1 = np.random.random([4, 6, 6]).astype("int32")
            # input is a variable which shape is [4, 6, 6]
H
hong 已提交
570
            input = paddle.to_tensor(input_1)
571 572 573 574 575
            x0, x1, x2 = paddle.split(input, num_or_sections=3, axis=1)
            x0_out = x0.numpy()
            x1_out = x1.numpy()
            x2_out = x2.numpy()
            ex_x0, ex_x1, ex_x2 = np.split(input_1, 3, axis=1)
H
hong 已提交
576

577 578 579 580 581 582 583 584 585 586 587 588 589 590 591
            # input is a variable which shape is [4, 6, 6]
            input = paddle.to_tensor(input_1)
            input.stop_gradient = False
            x0, x1, x2 = paddle.split(input, num_or_sections=3, axis=1)
            eager_x0_out = x0.numpy()
            eager_x1_out = x1.numpy()
            eager_x2_out = x2.numpy()
            loss = x0.sum()
            loss.backward()
            manul_grad = np.zeros_like(input_1)
            manul_grad[:, :2, :] = 1
            np.testing.assert_allclose(input.gradient(), manul_grad, rtol=1e-05)
            np.testing.assert_allclose(ex_x0, eager_x0_out, rtol=1e-05)
            np.testing.assert_allclose(ex_x1, eager_x1_out, rtol=1e-05)
            np.testing.assert_allclose(ex_x2, eager_x2_out, rtol=1e-05)
H
hong 已提交
592

593 594 595
        np.testing.assert_allclose(ex_x0, x0_out, rtol=1e-05)
        np.testing.assert_allclose(ex_x1, x1_out, rtol=1e-05)
        np.testing.assert_allclose(ex_x2, x2_out, rtol=1e-05)
596 597 598 599 600

    def test_out2(self):
        with fluid.dygraph.guard():
            input_1 = np.random.random([4, 6, 6]).astype("bool")
            # input is a variable which shape is [4, 6, 6]
H
hong 已提交
601
            input = paddle.to_tensor(input_1)
602 603 604 605 606
            x0, x1, x2 = paddle.split(input, num_or_sections=3, axis=1)
            x0_out = x0.numpy()
            x1_out = x1.numpy()
            x2_out = x2.numpy()
            ex_x0, ex_x1, ex_x2 = np.split(input_1, 3, axis=1)
607 608 609
        np.testing.assert_allclose(ex_x0, x0_out, rtol=1e-05)
        np.testing.assert_allclose(ex_x1, x1_out, rtol=1e-05)
        np.testing.assert_allclose(ex_x2, x2_out, rtol=1e-05)
610

C
Charles-hit 已提交
611 612 613 614 615 616 617 618 619 620 621
    def test_out3(self):
        with fluid.dygraph.guard():
            np.random.seed(2021)
            input_1 = np.random.random([4, 6, 6]).astype("int32")
            # input is a variable which shape is [4, 6, 6]
            input = paddle.to_tensor(input_1)
            out_dy = paddle.split(input, [6], axis=1)
            out_dy = out_dy[0]
            out_dy_np = out_dy.numpy()
            ex_out = np.split(input_1, [6], axis=1)
            ex_out = ex_out[0]
622 623 624 625 626
            input = paddle.to_tensor(input_1)
            out_eager = paddle.split(input, [6], axis=1)
            out_eager = out_eager[0]
            out_eager_np = out_dy.numpy()
            np.testing.assert_allclose(ex_out, out_eager_np, rtol=1e-05)
C
Charles-hit 已提交
627 628
        np.testing.assert_allclose(ex_out, out_dy_np, rtol=1e-05)

629 630 631 632
    def test_out_tensor_input(self):
        with fluid.dygraph.guard():
            input_1 = np.random.random([4, 6, 6]).astype("int32")
            # input is a variable which shape is [4, 6, 6]
H
hong 已提交
633
            input = paddle.to_tensor(input_1)
634
            num1 = paddle.full(shape=[1], fill_value=2, dtype='int32')
635 636 637
            x0, x1, x2 = paddle.split(
                input, num_or_sections=[num1, 2, 2], axis=1
            )
638 639 640 641
            x0_out = x0.numpy()
            x1_out = x1.numpy()
            x2_out = x2.numpy()
            ex_x0, ex_x1, ex_x2 = np.split(input_1, 3, axis=1)
642 643 644
        np.testing.assert_allclose(ex_x0, x0_out, rtol=1e-05)
        np.testing.assert_allclose(ex_x1, x1_out, rtol=1e-05)
        np.testing.assert_allclose(ex_x2, x2_out, rtol=1e-05)
645 646

    def test_axis_tensor_input(self):
647 648 649
        with fluid.dygraph.guard():
            input_1 = np.random.random([4, 6, 6]).astype("int32")
            # input is a variable which shape is [4, 6, 6]
H
hong 已提交
650
            input = paddle.to_tensor(input_1)
651
            num1 = paddle.full(shape=[1], fill_value=1, dtype='int32')
652 653 654
            x0, x1, x2 = paddle.split(
                input, num_or_sections=[2, 2, 2], axis=num1
            )
655 656 657 658
            x0_out = x0.numpy()
            x1_out = x1.numpy()
            x2_out = x2.numpy()
            ex_x0, ex_x1, ex_x2 = np.split(input_1, 3, axis=1)
659 660 661
        np.testing.assert_allclose(ex_x0, x0_out, rtol=1e-05)
        np.testing.assert_allclose(ex_x1, x1_out, rtol=1e-05)
        np.testing.assert_allclose(ex_x2, x2_out, rtol=1e-05)
662

663
    def test_negative_one_section(self):
664 665 666 667 668 669 670
        with fluid.dygraph.guard():
            input_1 = np.random.random([4, 6, 6]).astype("int32")
            # input is a variable which shape is [4, 6, 6]
            input = paddle.to_tensor(input_1)
            num1 = paddle.full(shape=[1], fill_value=1, dtype='int32')
            x0 = paddle.split(input, num_or_sections=[-1], axis=num1)
            x0_out = x0[0].numpy()
671
        np.testing.assert_array_equal(x0_out, input.numpy())
672

673

674 675 676 677 678 679 680 681 682 683
class API_TestEmptySplit(unittest.TestCase):
    def test_axis_input_empty_section(self):
        with fluid.dygraph.guard():
            input_1 = np.random.random([8, 6, 6]).astype("float32")
            # input is a variable which shape is [8, 6, 6]
            input = paddle.to_tensor(input_1)
            x0, x1, x2 = paddle.split(input, num_or_sections=[5, 0, 3])
            x0_out = x0.numpy()
            x1_out = x1.numpy()
            x2_out = x2.numpy()
684 685 686 687 688 689 690
            ex_x0, ex_x1, ex_x2 = np.split(
                input_1,
                [
                    5,
                    5,
                ],
            )
691 692 693
        np.testing.assert_allclose(ex_x0, x0_out, rtol=1e-05)
        np.testing.assert_allclose(ex_x1, x1_out, rtol=1e-05)
        np.testing.assert_allclose(ex_x2, x2_out, rtol=1e-05)
694 695


Y
Yancey 已提交
696
if __name__ == '__main__':
697
    paddle.enable_static()
Y
Yancey 已提交
698
    unittest.main()