1import math
2import warnings
3
4import numpy
5
6import chainer
7from chainer.backends import cuda
8from chainer import function_node
9from chainer import utils
10from chainer.utils import type_check
11
12
13_erf_cpu = None
14
15
16class Erf(function_node.FunctionNode):
17
18    @property
19    def label(self):
20        return 'erf'
21
22    def check_type_forward(self, in_types):
23        type_check._argname(in_types, ('x',))
24        type_check.expect(in_types[0].dtype.kind == 'f')
25
26    def forward_cpu(self, x):
27        global _erf_cpu
28        if _erf_cpu is None:
29            try:
30                from scipy import special
31                _erf_cpu = special.erf
32            except ImportError:
33                warnings.warn(
34                    'SciPy is not available. Forward computation of erf in CPU'
35                    ' can be slow without SciPy.',
36                    chainer.warnings.PerformanceWarning)
37                _erf_cpu = numpy.vectorize(math.erf)
38        self.retain_inputs((0,))
39        return utils.force_array(_erf_cpu(x[0]), dtype=x[0].dtype),
40
41    def forward_gpu(self, x):
42        self.retain_inputs((0,))
43        return cuda.elementwise(
44            'T x', 'T y',
45            'y = erf(x)',
46            'elementwise_erf',
47        )(x[0]),
48
49    def backward(self, indexes, gy):
50        x = self.get_retained_inputs()[0]
51        return 2 / numpy.pi ** 0.5 * chainer.functions.exp(-x ** 2) * gy[0],
52
53
54def erf(x):
55    """Elementwise error function.
56
57    .. note::
58       Forward computation in CPU can be slow if
59       `SciPy <https://www.scipy.org/>`_ is not available.
60
61    Args:
62        x (:class:`~chainer.Variable` or :ref:`ndarray`): Input variable.
63
64    Returns:
65        ~chainer.Variable: Output variable.
66    """
67    return Erf().apply((x,))[0]
68