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


@testing.parameterize(*testing.product({
    'shape': [(3, 4), ()],
    'dtype': [numpy.float16, numpy.float32, numpy.float64],
}))
class TestFlatten(unittest.TestCase):

    dtype = numpy.float32

    def setUp(self):
        self.x = numpy.random.uniform(-1, 1, self.shape).astype(self.dtype)
        self.g_shape = (numpy.prod((1,) + self.shape),)
        self.g = numpy.random.uniform(-1, 1, self.g_shape).astype(self.dtype)

    def check_forward(self, x_data):
        x = chainer.Variable(x_data)
        y = functions.flatten(x)

        self.assertEqual(y.shape, self.g_shape)
        self.assertEqual(y.dtype, self.dtype)
        testing.assert_allclose(self.x.flatten(), y.data)

    def test_forward_cpu(self):
        self.check_forward(self.x)

    @attr.gpu
    def test_forward_gpu(self):
        self.check_forward(cuda.to_gpu(self.x))

    def check_backward(self, x_data, g_data):
        gradient_check.check_backward(
            functions.Flatten(), x_data, g_data, dtype=numpy.float64)

    def test_backward_cpu(self):
        self.check_backward(self.x, self.g)

    @attr.gpu
    def test_backward_gpu(self):
        self.check_backward(cuda.to_gpu(self.x), cuda.to_gpu(self.g))


testing.run_module(__name__, __file__)

[(53, 'chainer.testing.run_module', 'testing.run_module', False, 'from chainer import testing\n'), (27, 'chainer.Variable', 'chainer.Variable', False, 'import chainer\n'), (28, 'chainer.functions.flatten', 'functions.flatten', False, 'from chainer import functions\n'), (13, 'chainer.testing.product', 'testing.product', False, 'from chainer import testing\n'), (23, 'numpy.prod', 'numpy.prod', False, 'import numpy\n'), (39, 'chainer.cuda.to_gpu', 'cuda.to_gpu', False, 'from chainer import cuda\n'), (43, 'chainer.functions.Flatten', 'functions.Flatten', False, 'from chainer import functions\n'), (50, 'chainer.cuda.to_gpu', 'cuda.to_gpu', False, 'from chainer import cuda\n'), (50, 'chainer.cuda.to_gpu', 'cuda.to_gpu', False, 'from chainer import cuda\n'), (22, 'numpy.random.uniform', 'numpy.random.uniform', False, 'import numpy\n'), (24, 'numpy.random.uniform', 'numpy.random.uniform', False, 'import numpy\n')]