1import numpy
2
3try:
4    from scipy import special
5    available_cpu = True
6except ImportError as e:
7    available_cpu = False
8    _import_error = e
9
10from chainer.backends import cuda
11from chainer import function_node
12from chainer import utils
13from chainer.utils import type_check
14
15
16class Erfcx(function_node.FunctionNode):
17
18    @property
19    def label(self):
20        return 'erfcx'
21
22    def check_type_forward(self, in_types):
23        type_check.expect(in_types.size() == 1)
24        type_check.expect(in_types[0].dtype.kind == 'f')
25
26    def forward_cpu(self, x):
27        if not available_cpu:
28            raise ImportError('SciPy is not available. Forward computation'
29                              ' of erfcx in CPU cannot be done. ' +
30                              str(_import_error))
31        self.retain_inputs((0,))
32        self.retain_outputs((0,))
33        return utils.force_array(special.erfcx(x[0]), dtype=x[0].dtype),
34
35    def forward_gpu(self, x):
36        self.retain_inputs((0,))
37        self.retain_outputs((0,))
38        return cuda.elementwise(
39            'T x', 'T y',
40            'y = erfcx(x)',
41            'elementwise_erfcx',
42        )(x[0]),
43
44    def backward(self, indexes, gy):
45        x = self.get_retained_inputs()[0]
46        y = self.get_retained_outputs()[0]
47        return 2 * (x * y - numpy.pi ** -0.5) * gy[0],
48
49
50def erfcx(x):
51    """Elementwise scaled complementary error function.
52
53    .. note::
54       Forward computation in CPU cannot be done if
55       `SciPy <https://www.scipy.org/>`_ is not available.
56
57    Args:
58        x (:class:`~chainer.Variable` or :ref:`ndarray`): Input variable.
59
60    Returns:
61        ~chainer.Variable: Output variable.
62    """
63    return Erfcx().apply((x,))[0]
64