<commit_before>import numpy

import chainer
from chainer import cuda
from chainer import function_node
from chainer.functions.array import broadcast
from chainer.functions.math import sum as sum_
from chainer.utils import type_check

if cuda.cudnn_enabled:
    cudnn = cuda.cudnn
    libcudnn = cudnn.cudnn
    _algorithm = libcudnn.CUDNN_SOFTMAX_ACCURATE
    _mode = libcudnn.CUDNN_SOFTMAX_MODE_CHANNEL


def _get_tensor4d_shape(axis, shape):
    left_shape = numpy.prod(shape[slice(0, axis)], dtype=numpy.int)
    center_shape = shape[axis]
    right_shape = numpy.prod(
        shape[slice(axis + 1, len(shape))], dtype=numpy.int)
    return left_shape, center_shape, right_shape, 1


class Softmax(function_node.FunctionNode):

    """Softmax activation function."""

    def __init__(self, axis=1):
        self.axis = axis

    def check_type_forward(self, in_types):
        type_check.expect(in_types.size() == 1)
        x_type, = in_types

        type_check.expect(
            x_type.dtype.kind == 'f',
            x_type.ndim > 1,
            self.axis < x_type.ndim
        )

    def forward(self, x):
        xp = cuda.get_array_module(*x)
        if xp is not numpy and chainer.should_use_cudnn('>=auto'):
            oz_dtype = 'd' if x[0].dtype == 'd' else 'f'
            one = numpy.array(1, dtype=oz_dtype).ctypes
            zero = numpy.array(0, dtype=oz_dtype).ctypes
            handle = cudnn.get_handle()
            x_tensor4d = cuda.cupy.ascontiguousarray(
                x[0].reshape(_get_tensor4d_shape(self.axis, x[0].shape)))
            desc = cudnn.create_tensor_descriptor(x_tensor4d)
            y = xp.empty_like(x[0])
            libcudnn.softmaxForward(
                handle, _algorithm, _mode, one.data, desc.value,
                x_tensor4d.data.ptr, zero.data, desc.value,
                y.data.ptr)
        else:
            y = x[0] - x[0].max(axis=self.axis, keepdims=True)
            xp.exp(y, out=y)
            y /= y.sum(axis=self.axis, keepdims=True)

        self.retain_outputs((0,))
        return y,

    def backward(self, indexes, grad_outputs):
        y = self.get_retained_outputs()[0]
        gy, = grad_outputs
        return _SoftmaxGrad(self.axis).apply((y, gy))


class _SoftmaxGrad(function_node.FunctionNode):

    def __init__(self, axis):
        self.axis = axis

    def forward(self, inputs):
        self.retain_inputs((0, 1))
        y, gy = inputs
        xp = cuda.get_array_module(*y)
        if xp is not numpy and chainer.should_use_cudnn('>=auto'):
            oz_dtype = 'd' if y[0].dtype == 'd' else 'f'
            one = numpy.array(1, dtype=oz_dtype).ctypes
            zero = numpy.array(0, dtype=oz_dtype).ctypes
            handle = cudnn.get_handle()
            gx = xp.empty_like(y)
            gx_tensor4d = cuda.cupy.ascontiguousarray(
                gx.reshape(_get_tensor4d_shape(self.axis, gx.shape)))
            gy = cuda.cupy.ascontiguousarray(gy)
            desc = cudnn.create_tensor_descriptor(gx_tensor4d)
            libcudnn.softmaxBackward(
                handle, _algorithm, _mode, one.data, desc.value,
                y.data.ptr, desc.value, gy.data.ptr, zero.data,
                desc.value, gx.data.ptr)
        else:
            gx = y * gy
            sumdx = gx.sum(axis=self.axis, keepdims=True)
            gx -= y * sumdx

        return gx,

    def backward(self, indexes, grad_outputs):
        y, gy = self.get_retained_inputs()
        ggx, = grad_outputs
        gs = sum_.sum(-ggx * y, axis=self.axis, keepdims=True)
        ga = ggx + broadcast.broadcast_to(gs, gy.shape)
        ret = []
        if 0 in indexes:
            s = broadcast.broadcast_to(
                sum_.sum(y * gy, axis=self.axis, keepdims=True), gy.shape)
            gy2 = -ggx * s + ga * gy
            ret.append(gy2)
        if 1 in indexes:
            ggy = ga * y
            ret.append(ggy)
        return tuple(ret)


def softmax(x, axis=1):
    """Softmax function.

    This function computes its softmax along an axis. Let
    :math:`c = (c_1, c_2, \\dots, c_D)` be the slice of ``x`` along with
    the axis. For each slice :math:`c`, it computes the function :math:`f(c)`
    defined as :math:`f(c)={\\exp(c) \\over \\sum_{d} \\exp(c_d)}`.

    Args:
        x (:class:`~chainer.Variable` or :class:`numpy.ndarray` or \
        :class:`cupy.ndarray`):
            Input variable.
            A :math:`n`-dimensional (:math:`n \\geq 2`) float array.
        axis (int): The axis along which the softmax is to be computed.

    Returns:
        ~chainer.Variable: Output variable.
        A :math:`n`-dimensional (:math:`n \\geq 2`) float array, which is the
        same shape with x.

    .. admonition:: Example

        >>> x = np.array([[0, 1, 2], [0, 2, 4]], 'f')
        >>> x
        array([[ 0.,  1.,  2.],
               [ 0.,  2.,  4.]], dtype=float32)
        >>> y = F.softmax(x, axis=1)
        >>> y.data
        array([[ 0.09003057,  0.24472848,  0.66524094],
               [ 0.01587624,  0.11731043,  0.86681336]], dtype=float32)
        >>> F.sum(y, axis=1).data
        array([ 1.,  1.], dtype=float32)

    """
    return Softmax(axis=axis).apply((x,))[0]
<commit_msg>Use chainer.functions namespace.<commit_after>import numpy

import chainer
from chainer import cuda
from chainer import function_node
import chainer.functions
from chainer.utils import type_check

if cuda.cudnn_enabled:
    cudnn = cuda.cudnn
    libcudnn = cudnn.cudnn
    _algorithm = libcudnn.CUDNN_SOFTMAX_ACCURATE
    _mode = libcudnn.CUDNN_SOFTMAX_MODE_CHANNEL


def _get_tensor4d_shape(axis, shape):
    left_shape = numpy.prod(shape[slice(0, axis)], dtype=numpy.int)
    center_shape = shape[axis]
    right_shape = numpy.prod(
        shape[slice(axis + 1, len(shape))], dtype=numpy.int)
    return left_shape, center_shape, right_shape, 1


class Softmax(function_node.FunctionNode):

    """Softmax activation function."""

    def __init__(self, axis=1):
        self.axis = axis

    def check_type_forward(self, in_types):
        type_check.expect(in_types.size() == 1)
        x_type, = in_types

        type_check.expect(
            x_type.dtype.kind == 'f',
            x_type.ndim > 1,
            self.axis < x_type.ndim
        )

    def forward(self, x):
        xp = cuda.get_array_module(*x)
        if xp is not numpy and chainer.should_use_cudnn('>=auto'):
            oz_dtype = 'd' if x[0].dtype == 'd' else 'f'
            one = numpy.array(1, dtype=oz_dtype).ctypes
            zero = numpy.array(0, dtype=oz_dtype).ctypes
            handle = cudnn.get_handle()
            x_tensor4d = cuda.cupy.ascontiguousarray(
                x[0].reshape(_get_tensor4d_shape(self.axis, x[0].shape)))
            desc = cudnn.create_tensor_descriptor(x_tensor4d)
            y = xp.empty_like(x[0])
            libcudnn.softmaxForward(
                handle, _algorithm, _mode, one.data, desc.value,
                x_tensor4d.data.ptr, zero.data, desc.value,
                y.data.ptr)
        else:
            y = x[0] - x[0].max(axis=self.axis, keepdims=True)
            xp.exp(y, out=y)
            y /= y.sum(axis=self.axis, keepdims=True)

        self.retain_outputs((0,))
        return y,

    def backward(self, indexes, grad_outputs):
        y = self.get_retained_outputs()[0]
        gy, = grad_outputs
        return _SoftmaxGrad(self.axis).apply((y, gy))


class _SoftmaxGrad(function_node.FunctionNode):

    def __init__(self, axis):
        self.axis = axis

    def forward(self, inputs):
        self.retain_inputs((0, 1))
        y, gy = inputs
        xp = cuda.get_array_module(*y)
        if xp is not numpy and chainer.should_use_cudnn('>=auto'):
            oz_dtype = 'd' if y[0].dtype == 'd' else 'f'
            one = numpy.array(1, dtype=oz_dtype).ctypes
            zero = numpy.array(0, dtype=oz_dtype).ctypes
            handle = cudnn.get_handle()
            gx = xp.empty_like(y)
            gx_tensor4d = cuda.cupy.ascontiguousarray(
                gx.reshape(_get_tensor4d_shape(self.axis, gx.shape)))
            gy = cuda.cupy.ascontiguousarray(gy)
            desc = cudnn.create_tensor_descriptor(gx_tensor4d)
            libcudnn.softmaxBackward(
                handle, _algorithm, _mode, one.data, desc.value,
                y.data.ptr, desc.value, gy.data.ptr, zero.data,
                desc.value, gx.data.ptr)
        else:
            gx = y * gy
            sumdx = gx.sum(axis=self.axis, keepdims=True)
            gx -= y * sumdx

        return gx,

    def backward(self, indexes, grad_outputs):
        y, gy = self.get_retained_inputs()
        ggx, = grad_outputs
        gs = chainer.functions.sum(-ggx * y, axis=self.axis, keepdims=True)
        ga = ggx + chainer.functions.broadcast_to(gs, gy.shape)
        ret = []
        if 0 in indexes:
            s = chainer.functions.broadcast_to(chainer.functions.sum(
                y * gy, axis=self.axis, keepdims=True), gy.shape)
            gy2 = -ggx * s + ga * gy
            ret.append(gy2)
        if 1 in indexes:
            ggy = ga * y
            ret.append(ggy)
        return tuple(ret)


def softmax(x, axis=1):
    """Softmax function.

    This function computes its softmax along an axis. Let
    :math:`c = (c_1, c_2, \\dots, c_D)` be the slice of ``x`` along with
    the axis. For each slice :math:`c`, it computes the function :math:`f(c)`
    defined as :math:`f(c)={\\exp(c) \\over \\sum_{d} \\exp(c_d)}`.

    Args:
        x (:class:`~chainer.Variable` or :class:`numpy.ndarray` or \
        :class:`cupy.ndarray`):
            Input variable.
            A :math:`n`-dimensional (:math:`n \\geq 2`) float array.
        axis (int): The axis along which the softmax is to be computed.

    Returns:
        ~chainer.Variable: Output variable.
        A :math:`n`-dimensional (:math:`n \\geq 2`) float array, which is the
        same shape with x.

    .. admonition:: Example

        >>> x = np.array([[0, 1, 2], [0, 2, 4]], 'f')
        >>> x
        array([[ 0.,  1.,  2.],
               [ 0.,  2.,  4.]], dtype=float32)
        >>> y = F.softmax(x, axis=1)
        >>> y.data
        array([[ 0.09003057,  0.24472848,  0.66524094],
               [ 0.01587624,  0.11731043,  0.86681336]], dtype=float32)
        >>> F.sum(y, axis=1).data
        array([ 1.,  1.], dtype=float32)

    """
    return Softmax(axis=axis).apply((x,))[0]

{'fraction_non_alphanumeric': 0.1318797425799635, 'fraction_numerical': 0.022092018057823456, 'mean_word_length': 2.0914489311163895, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 39, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2542997', 'n_tokens_mistral': 3741, 'n_tokens_neox': 3576, 'n_words': 786}