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, 0, 53, 38), 'chainer.testing.run_module', 'testing.run_module', ({(53, 19, 53, 27): '__name__', (53, 29, 53, 37): '__file__'}, {}), '(__name__, __file__)', False, 'from chainer import testing\n'), ((27, 12, 27, 36), 'chainer.Variable', 'chainer.Variable', ({(27, 29, 27, 35): 'x_data'}, {}), '(x_data)', False, 'import chainer\n'), ((28, 12, 28, 32), 'chainer.functions.flatten', 'functions.flatten', ({(28, 30, 28, 31): 'x'}, {}), '(x)', False, 'from chainer import functions\n'), ((13, 23, 16, 2), 'chainer.testing.product', 'testing.product', ({(13, 39, 16, 1): "{'shape': [(3, 4), ()], 'dtype': [numpy.float16, numpy.float32, numpy.float64]}"}, {}), "({'shape': [(3, 4), ()], 'dtype': [numpy.float16, numpy.\n    float32, numpy.float64]})", False, 'from chainer import testing\n'), ((23, 24, 23, 53), 'numpy.prod', 'numpy.prod', ({(23, 35, 23, 52): '((1,) + self.shape)'}, {}), '((1,) + self.shape)', False, 'import numpy\n'), ((39, 27, 39, 46), 'chainer.cuda.to_gpu', 'cuda.to_gpu', ({(39, 39, 39, 45): 'self.x'}, {}), '(self.x)', False, 'from chainer import cuda\n'), ((43, 12, 43, 31), 'chainer.functions.Flatten', 'functions.Flatten', ({}, {}), '()', False, 'from chainer import functions\n'), ((50, 28, 50, 47), 'chainer.cuda.to_gpu', 'cuda.to_gpu', ({(50, 40, 50, 46): 'self.x'}, {}), '(self.x)', False, 'from chainer import cuda\n'), ((50, 49, 50, 68), 'chainer.cuda.to_gpu', 'cuda.to_gpu', ({(50, 61, 50, 67): 'self.g'}, {}), '(self.g)', False, 'from chainer import cuda\n'), ((22, 17, 22, 56), 'numpy.random.uniform', 'numpy.random.uniform', ({(22, 38, 22, 40): '-1', (22, 42, 22, 43): '1', (22, 45, 22, 55): 'self.shape'}, {}), '(-1, 1, self.shape)', False, 'import numpy\n'), ((24, 17, 24, 58), 'numpy.random.uniform', 'numpy.random.uniform', ({(24, 38, 24, 40): '-1', (24, 42, 24, 43): '1', (24, 45, 24, 57): 'self.g_shape'}, {}), '(-1, 1, self.g_shape)', False, 'import numpy\n')]