import unittest

import numpy

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


class TestScale(unittest.TestCase):

    def setUp(self):
        self.x1 = numpy.random.uniform(-1, 1, (3, 2, 3)).astype(numpy.float32)
        self.x2 = numpy.random.uniform(-1, 1, (2)).astype(numpy.float32)
        self.axis = 1
        self.y_expected = numpy.copy(self.x1)
        for i, j, k in numpy.ndindex(self.y_expected.shape):
            self.y_expected[i, j, k] *= self.x2[j]
        self.gy = numpy.random.uniform(-1, 1, (3, 2, 3)).astype(numpy.float32)

        self.check_backward_options = {'atol': 1e-2, 'rtol': 1e-2}

    def check_forward(self, x1_data, x2_data, axis, y_expected):
        x1 = chainer.Variable(x1_data)
        x2 = chainer.Variable(x2_data)
        y = functions.scale(x1, x2, axis)
        testing.assert_allclose(y_expected, y.data)

    def test_forward_cpu(self):
        self.check_forward(self.x1, self.x2, self.axis, self.y_expected)

    @attr.gpu
    def test_forward_gpu(self):
        x1 = cuda.to_gpu(self.x1)
        x2 = cuda.to_gpu(self.x2)
        self.check_forward(x1, x2, self.axis, self.y_expected)

    def check_backward(self, x1_data, x2_data, axis, y_grad):
        x = (x1_data, x2_data)
        gradient_check.check_backward(
            lambda x, y: functions.scale(x, y, axis),
            x, y_grad,
            **self.check_backward_options)

    def test_backward_cpu(self):
        self.check_backward(self.x1, self.x2, self.axis, self.gy)

    @attr.gpu
    def test_backward_gpu(self):
        x1 = cuda.to_gpu(self.x1)
        x2 = cuda.to_gpu(self.x2)
        gy = cuda.to_gpu(self.gy)
        self.check_backward(x1, x2, self.axis, gy)


class TestScaleInvalidShape(unittest.TestCase):

    def test_scale_invalid_shape(self):
        x1 = chainer.Variable(numpy.zeros((3, 2, 3), numpy.float32))
        x2 = chainer.Variable(numpy.zeros((2), numpy.float32))
        axis = 0
        with chainer.using_config('debug', True):
            with self.assertRaises(AssertionError):
                functions.scale(x1, x2, axis)


testing.run_module(__name__, __file__)

import unittest

import numpy

import chainer
from chainer import functions
from chainer import testing


@testing.inject_backend_tests(
    None,
    # CPU tests
    [
        {},
    ]
    # GPU tests
    + testing.product({
        'use_cuda': [True],
        'cuda_device': [0, 1],
    })
    # ChainerX tests
    + testing.product({
        'use_chainerx': [True],
        'chainerx_device': ['native:0', 'cuda:0', 'cuda:1']
    })
)
class TestScale(testing.FunctionTestCase):

    def setUp(self):
        self.axis = 1
        self.check_backward_options.update({'atol': 1e-5, 'rtol': 1e-5})
        self.check_double_backward_options.update(
            {'atol': 1e-4, 'rtol': 1e-4})

    def generate_inputs(self):
        x1 = numpy.random.uniform(-1, 1, (3, 2, 3)).astype(numpy.float32)
        x2 = numpy.random.uniform(-1, 1, (2)).astype(numpy.float32)
        return x1, x2

    def forward_expected(self, inputs):
        x1, x2 = inputs
        y_expected = numpy.copy(x1)
        for i, j, k in numpy.ndindex(y_expected.shape):
            y_expected[i, j, k] *= x2[j]
        return y_expected,

    def forward(self, inputs, device):
        x1, x2 = inputs
        y = functions.scale(x1, x2, self.axis)
        return y,


class TestScaleInvalidShape(unittest.TestCase):

    def test_scale_invalid_shape(self):
        x1 = chainer.Variable(numpy.zeros((3, 2, 3), numpy.float32))
        x2 = chainer.Variable(numpy.zeros((2), numpy.float32))
        axis = 0
        with chainer.using_config('debug', True):
            with self.assertRaises(AssertionError):
                functions.scale(x1, x2, axis)


testing.run_module(__name__, __file__)

Simplify `F.scale` test