import unittest

import numpy

import chainer
from chainer import cuda
from chainer import functions
from chainer import gradient_check
from chainer import testing
from chainer.testing import attr


class ConcatTestBase(object):

    def check_forward(self, xs_data, y_data, axis):
        xs = tuple(chainer.Variable(x_data) for x_data in xs_data)
        y = functions.concat(xs, axis=axis)
        gradient_check.assert_allclose(y_data, y.data, atol=0, rtol=0)
        self.assertIsInstance(y.data.shape, tuple)

    def test_forward_cpu(self):
        self.check_forward(self.xs, self.y, axis=self.axis)

    @attr.gpu
    def test_forward_gpu(self):
        self.check_forward(
            [cuda.to_gpu(x.copy()) for x in self.xs],
            cuda.to_gpu(self.y), axis=self.axis)

    def check_backward(self, xs_data, axis):
        xs = tuple(chainer.Variable(x_data) for x_data in xs_data)
        y = functions.concat(xs, axis=axis)
        y.grad = y.data
        y.backward()

        for x in xs:
            gradient_check.assert_allclose(x.data, x.grad, atol=0, rtol=0)

    def test_backward_cpu(self):
        self.check_backward(self.xs, axis=self.axis)

    @attr.gpu
    def test_backward_gpu(self):
        self.check_backward([cuda.to_gpu(x.copy()) for x in self.xs],
                            axis=self.axis)


class TestConcat1(unittest.TestCase, ConcatTestBase):

    def setUp(self):
        self.y = numpy.arange(42, dtype=numpy.float32).reshape(2, 7, 3)
        self.xs = [self.y[:, :2], self.y[:, 2:5], self.y[:, 5:]]
        self.axis = 1


class TestConcat2(unittest.TestCase, ConcatTestBase):

    def setUp(self):
        self.y = numpy.arange(21, dtype=numpy.float32).reshape(7, 3)
        self.xs = [self.y[:2], self.y[2:5], self.y[5:]]
        self.axis = 0


class TestConcatLastAxis(unittest.TestCase, ConcatTestBase):

    def setUp(self):
        self.y = numpy.arange(2, dtype=numpy.float32)
        self.xs = [self.y[:1], self.y[1:]]
        self.axis = 0


class TestConcatOneElement(unittest.TestCase, ConcatTestBase):

    def setUp(self):
        self.y = numpy.arange(2, dtype=numpy.float32)
        self.xs = [self.y]
        self.axis = 0


testing.run_module(__name__, __file__)

@@ -196,37 +196,630 @@
 r%0A%0A%0A
-class ConcatTestBase(object):
+@testing.parameterize(%0A    %7B'shape': (2, 7, 3), 'axis': 1,%0A     'slices': %5B%5Bslice(None), slice(None, 2)%5D, %5Bslice(None), slice(2, 5)%5D,%0A                %5Bslice(None), slice(5, None)%5D%5D%7D,%0A    %7B'shape': (7, 3), 'axis': 0,%0A     'slices': %5Bslice(None, 2), slice(2, 5), slice(5, None)%5D%7D,%0A    %7B'shape': (2,), 'axis': 0, 'slices': %5Bslice(None, 1), slice(1, None)%5D%7D,%0A    %7B'shape': (2,), 'axis': 0, 'slices': %5B()%5D%7D,%0A)%0Aclass ConcatTest(unittest.TestCase):%0A%0A    def setUp(self):%0A        self.y = numpy.arange(%0A            numpy.prod(self.shape), dtype=numpy.float32).reshape(self.shape)%0A        self.xs = %5Bself.y%5Bs%5D for s in self.slices%5D
 %0A%0A  
@@ -1926,864 +1926,8 @@
 )%0A%0A%0A
-class TestConcat1(unittest.TestCase, ConcatTestBase):%0A%0A    def setUp(self):%0A        self.y = numpy.arange(42, dtype=numpy.float32).reshape(2, 7, 3)%0A        self.xs = %5Bself.y%5B:, :2%5D, self.y%5B:, 2:5%5D, self.y%5B:, 5:%5D%5D%0A        self.axis = 1%0A%0A%0Aclass TestConcat2(unittest.TestCase, ConcatTestBase):%0A%0A    def setUp(self):%0A        self.y = numpy.arange(21, dtype=numpy.float32).reshape(7, 3)%0A        self.xs = %5Bself.y%5B:2%5D, self.y%5B2:5%5D, self.y%5B5:%5D%5D%0A        self.axis = 0%0A%0A%0Aclass TestConcatLastAxis(unittest.TestCase, ConcatTestBase):%0A%0A    def setUp(self):%0A        self.y = numpy.arange(2, dtype=numpy.float32)%0A        self.xs = %5Bself.y%5B:1%5D, self.y%5B1:%5D%5D%0A        self.axis = 0%0A%0A%0Aclass TestConcatOneElement(unittest.TestCase, ConcatTestBase):%0A%0A    def setUp(self):%0A        self.y = numpy.arange(2, dtype=numpy.float32)%0A        self.xs = %5Bself.y%5D%0A        self.axis = 0%0A%0A%0A
 test

Simplify concat test