<commit_before>from chainer import cuda
from chainer import function_node
from chainer import functions
from chainer.utils import type_check


class LayerNormalization(function_node.FunctionNode):

    """Layer normalization"""

    def __init__(self, eps=1e-5):
        self.eps = eps

    def check_type_forward(self, in_types):
        type_check.expect(in_types.size() == 3)
        x_type, gamma_type, beta_type = in_types

        type_check.expect(
            x_type.dtype.kind == 'f',
            x_type.ndim == 2,
            gamma_type.ndim == 1,
            beta_type.ndim == 1,
            gamma_type.dtype == x_type.dtype,
            beta_type.dtype == x_type.dtype,
            gamma_type.shape == beta_type.shape,
        )

    def _compute(self, xp, x):
        mu = xp.mean(x, axis=1, keepdims=True)
        x_mu = x - xp.broadcast_to(mu, x.shape)
        squ_x_mu = xp.square(x_mu)
        var = xp.mean(squ_x_mu, axis=1, keepdims=True)
        std = xp.sqrt(var + self.eps)
        inv_std = 1. / std
        x_hat = x_mu * xp.broadcast_to(inv_std, x_mu.shape)
        return x_mu, var, inv_std, x_hat

    def forward(self, inputs):
        self.retain_inputs((0, 1))
        xp = cuda.get_array_module(*inputs)
        x, gamma, beta = inputs
        x_mu, var, inv_std, x_hat = self._compute(xp, x)
        scaled_x = x_hat * gamma[None, ]
        shifted_x = scaled_x + beta[None, ]
        return shifted_x,

    def backward(self, indexes, grad_outputs):
        F = functions
        x, gamma = self.get_retained_inputs()
        gy, = grad_outputs

        x_mu, var, inv_std, x_hat = self._compute(F, x)

        g_beta = F.sum(gy, axis=0)
        g_scaled_x = gy

        g_gamma = F.sum(g_scaled_x * x_hat, axis=0)<commit_msg>Restore _broadcast_to for NumPy 1.9<commit_after>from chainer import cuda
from chainer import function_node
from chainer import functions
from chainer.utils import type_check


def _broadcast_to(xp, x, shape):
    # xp: numpy, cupy, or chainer.functions
    if hasattr(xp, 'broadcast_to'):
        return xp.broadcast_to(x, shape)
    else:
        # numpy 1.9 doesn't support broadcast_to method
        dummy = xp.empty(shape)
        bx, _ = xp.broadcast_arrays(x, dummy)
        return bx


class LayerNormalization(function_node.FunctionNode):

    """Layer normalization"""

    def __init__(self, eps=1e-5):
        self.eps = eps

    def check_type_forward(self, in_types):
        type_check.expect(in_types.size() == 3)
        x_type, gamma_type, beta_type = in_types

        type_check.expect(
            x_type.dtype.kind == 'f',
            x_type.ndim == 2,
            gamma_type.ndim == 1,
            beta_type.ndim == 1,
            gamma_type.dtype == x_type.dtype,
            beta_type.dtype == x_type.dtype,
            gamma_type.shape == beta_type.shape,
        )

    def _compute(self, xp, x):
        mu = xp.mean(x, axis=1, keepdims=True)
        x_mu = x - _broadcast_to(xp, mu, x.shape)
        squ_x_mu = xp.square(x_mu)
        var = xp.mean(squ_x_mu, axis=1, keepdims=True)
        std = xp.sqrt(var + self.eps)
        inv_std = 1. / std
        x_hat = x_mu * _broadcast_to(xp, inv_std, x_mu.shape)
        return x_mu, var, inv_std, x_hat

    def forward(self, inputs):
        self.retain_inputs((0, 1))
        xp = cuda.get_array_module(*inputs)
        x, gamma, beta = inputs
        x_mu, var, inv_std, x_hat = self._compute(xp, x)
        scaled_x = x_hat * gamma[None, ]
        shifted_x = scaled_x + beta[None, ]
        return shifted_x,

    def backward(self, indexes, grad_outputs):
        F = functions
        x, gamma = self.get_retained_inputs()
        gy, = grad_outputs

        x_mu, var, inv_std, x_hat = self._compute(F, x)

        g_beta = F.sum(gy, axis=0)
        g_scaled_x = gy

        g_gamma = F.sum(g_scaled_x * x_hat, axis=0)
{'fraction_non_alphanumeric': 0.11108234075608493, 'fraction_numerical': 0.007767995857068877, 'mean_word_length': 2.2544229149115416, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, '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': '2137684', 'n_tokens_mistral': 1429, 'n_tokens_neox': 1377, 'n_words': 300}