<commit_before>import numpy

import chainer
from chainer import cuda
from chainer import function
from chainer.utils import type_check


if cuda.cudnn_enabled:
    cudnn = cuda.cudnn
    libcudnn = cuda.cudnn.cudnn


class PointerArray(object):

    def __init__(self, lst, back_pointer):
        self._value = numpy.array(lst, dtype=numpy.intp)
        # Store back_pointer to prevent the GC removes the original variable
        self._back_pointer = back_pointer

    @property
    def data(self):
        return self._value.ctypes.data


def _make_tensor_descriptor_array(xs):
    """Make an array of pointers denoting pointers of tensor descriptors. 
   """
    descs = []
    for x in xs:
        if x.ndim < 3:
            shape = x.shape + (1,) * (3 - x.ndim)
            x = x.reshape(shape)
        desc = cudnn.create_tensor_nd_descriptor(x)
        descs.append(desc)
    return PointerArray([d.value for d in descs], descs)


def _make_ptr_array(xs):
    """Make an array of pointers denoting pointers of ndarrays.
    """
    return PointerArray([x.data.ptr for x in xs], xs)


class DropoutStates(object):

    def __init__(self, states, desc):
        self.states = states
        self.desc = desc

    @staticmethod
    def create(handle, dropout, seed):
        states = cudnn.create_dropout_states(handle)
        desc = cudnn.create_dropout_descriptor(
            handle, dropout, states.data.ptr, states.size, seed)
        return DropoutStates(states, desc)

    @staticmethod
    def from_states(handle, states, dropout):
        desc = cudnn.create_dropout_descriptor(handle, dropout, 0, 0, 0)
        return DropoutStates(states, desc)


def _make_indices(batches):
    pos = 0
    inds = []
    for b in batches[:-1]:
        pos += b
        inds.append(pos)
    return inds


def _fetch_input(inputs, pos, n):
    return inputs[pos:pos+n], pos + n


class NStepLSTM(function.Function):

    def __init__(self, n_layers, states, train=True):
        self.n_layers = n_layers
        self.train = train
        self.states = states

    def check_type_forward(self, in_types):
        type_check.expect(in_types.size() > 2 + 16 * self.n_layers)
        h_type, c_type = in_types[:2]
        type_check.expect(
            h_type.dtype == numpy.float32,
            c_type.dtype == numpy.float32,

            h_type.ndim == 3,
            h_type.shape[0] == self.n_layers,
            c_type.ndim == 3,
            c_type.shape[0] == self.n_layers,

            # mini-batch size
            h_type.shape[1] == c_type.shape[1],

            # hidden size
            h_type.shape[2] == c_type.shape[2],
        )
        pos = 2
        w_types, pos = _fetch_input(in_types, pos, self.n_layers * 8)
        b_types, pos = _fetch_input(in_types, pos, self.n_layers * 8)
        x_types = in_types[pos:]
        for x_type in x_types:
            type_check.expect(
                x_type.dtype == numpy.float32,
                x_type.ndim == 2,
            )
        for x1_type, x2_type in zip(x_types, x_types[1:]):
            type_check.expect(
                # Check if xs are sorted by descending lengths
                x1_type.shape[0] >= x2_type.shape[0],
                x1_type.shape[1] == x2_type.shape[1])

        in_size = x_types[0].shape[1]
        out_size = h_type.shape[2]

        for layer in range(self.n_layers):
            for i in range(8):
                ind = layer * 8 + i
                w_type = w_types[ind]
                b_type = b_types[ind]
                if layer == 0 and i < 4:
                    w_in = in_size
                else:
                    w_in = out_size

                type_check.expect(
                    w_type.dtype == numpy.float32,
                    w_type.ndim == 2,
                    w_type.shape[0] == out_size,
                    w_type.shape[1] == w_in,

                    b_type.dtype == numpy.float32,
                    b_type.ndim == 1,
                    b_type.shape[0] == out_size,
                )

    def forward(self, inputs):
        (hx, cx), pos = _fetch_input(inputs, 0, 2)
        ws, pos = _fetch_input(inputs, pos, self.n_layers * 8)
        bs, pos = _fetch_input(inputs, pos, self.n_layers * 8)
        x_list = inputs[pos:]
        x = x_list[0]
        x = x.reshape(x.shape + (1,))
        x_desc = cudnn.create_tensor_nd_descriptor(x)

        length = len(x_list)
        n_units = hx.shape[2]

        xs = cuda.cupy.concatenate(x_list, axis=0)
        ys = cuda.cupy.empty((len(xs), n_units), dtype=xs.dtype)

        # shape of h and c is (batch_size, n_layer, hidden_size) in Chainer
        # but (hidden_size, batch_size, n_layer) in cuDNN
        handle = cudnn.get_handle()
        self.handle = handle

        rnn_desc = cudnn.create_rnn_descriptor(
            n_units, self.n_layers, self.states.desc,
            libcudnn.CUDNN_LINEAR_INPUT, libcudnn.CUDNN_UNIDIRECTIONAL,
            libcudnn.CUDNN_LSTM, libcudnn.CUDNN_DATA_FLOAT)
        self.rnn_desc = rnn_desc

        c_x_descs = _make_tensor_descriptor_array(x_list)
        hx_desc = cudnn.create_tensor_nd_descriptor(hx)
        cx_desc = cudnn.create_tensor_nd_descriptor(cx)

        weights_size = libcudnn.getRNNParamsSize(<commit_msg>Remove old comment<commit_after>import numpy

import chainer
from chainer import cuda
from chainer import function
from chainer.utils import type_check


if cuda.cudnn_enabled:
    cudnn = cuda.cudnn
    libcudnn = cuda.cudnn.cudnn


class PointerArray(object):

    def __init__(self, lst, back_pointer):
        self._value = numpy.array(lst, dtype=numpy.intp)
        # Store back_pointer to prevent the GC removes the original variable
        self._back_pointer = back_pointer

    @property
    def data(self):
        return self._value.ctypes.data


def _make_tensor_descriptor_array(xs):
    """Make an array of pointers denoting pointers of tensor descriptors. 
   """
    descs = []
    for x in xs:
        if x.ndim < 3:
            shape = x.shape + (1,) * (3 - x.ndim)
            x = x.reshape(shape)
        desc = cudnn.create_tensor_nd_descriptor(x)
        descs.append(desc)
    return PointerArray([d.value for d in descs], descs)


def _make_ptr_array(xs):
    """Make an array of pointers denoting pointers of ndarrays.
    """
    return PointerArray([x.data.ptr for x in xs], xs)


class DropoutStates(object):

    def __init__(self, states, desc):
        self.states = states
        self.desc = desc

    @staticmethod
    def create(handle, dropout, seed):
        states = cudnn.create_dropout_states(handle)
        desc = cudnn.create_dropout_descriptor(
            handle, dropout, states.data.ptr, states.size, seed)
        return DropoutStates(states, desc)

    @staticmethod
    def from_states(handle, states, dropout):
        desc = cudnn.create_dropout_descriptor(handle, dropout, 0, 0, 0)
        return DropoutStates(states, desc)


def _make_indices(batches):
    pos = 0
    inds = []
    for b in batches[:-1]:
        pos += b
        inds.append(pos)
    return inds


def _fetch_input(inputs, pos, n):
    return inputs[pos:pos+n], pos + n


class NStepLSTM(function.Function):

    def __init__(self, n_layers, states, train=True):
        self.n_layers = n_layers
        self.train = train
        self.states = states

    def check_type_forward(self, in_types):
        type_check.expect(in_types.size() > 2 + 16 * self.n_layers)
        h_type, c_type = in_types[:2]
        type_check.expect(
            h_type.dtype == numpy.float32,
            c_type.dtype == numpy.float32,

            h_type.ndim == 3,
            h_type.shape[0] == self.n_layers,
            c_type.ndim == 3,
            c_type.shape[0] == self.n_layers,

            # mini-batch size
            h_type.shape[1] == c_type.shape[1],

            # hidden size
            h_type.shape[2] == c_type.shape[2],
        )
        pos = 2
        w_types, pos = _fetch_input(in_types, pos, self.n_layers * 8)
        b_types, pos = _fetch_input(in_types, pos, self.n_layers * 8)
        x_types = in_types[pos:]
        for x_type in x_types:
            type_check.expect(
                x_type.dtype == numpy.float32,
                x_type.ndim == 2,
            )
        for x1_type, x2_type in zip(x_types, x_types[1:]):
            type_check.expect(
                # Check if xs are sorted by descending lengths
                x1_type.shape[0] >= x2_type.shape[0],
                x1_type.shape[1] == x2_type.shape[1])

        in_size = x_types[0].shape[1]
        out_size = h_type.shape[2]

        for layer in range(self.n_layers):
            for i in range(8):
                ind = layer * 8 + i
                w_type = w_types[ind]
                b_type = b_types[ind]
                if layer == 0 and i < 4:
                    w_in = in_size
                else:
                    w_in = out_size

                type_check.expect(
                    w_type.dtype == numpy.float32,
                    w_type.ndim == 2,
                    w_type.shape[0] == out_size,
                    w_type.shape[1] == w_in,

                    b_type.dtype == numpy.float32,
                    b_type.ndim == 1,
                    b_type.shape[0] == out_size,
                )

    def forward(self, inputs):
        (hx, cx), pos = _fetch_input(inputs, 0, 2)
        ws, pos = _fetch_input(inputs, pos, self.n_layers * 8)
        bs, pos = _fetch_input(inputs, pos, self.n_layers * 8)
        x_list = inputs[pos:]
        x = x_list[0]
        x = x.reshape(x.shape + (1,))
        x_desc = cudnn.create_tensor_nd_descriptor(x)

        length = len(x_list)
        n_units = hx.shape[2]

        xs = cuda.cupy.concatenate(x_list, axis=0)
        ys = cuda.cupy.empty((len(xs), n_units), dtype=xs.dtype)

        handle = cudnn.get_handle()
        self.handle = handle

        rnn_desc = cudnn.create_rnn_descriptor(
            n_units, self.n_layers, self.states.desc,
            libcudnn.CUDNN_LINEAR_INPUT, libcudnn.CUDNN_UNIDIRECTIONAL,
            libcudnn.CUDNN_LSTM, libcudnn.CUDNN_DATA_FLOAT)
        self.rnn_desc = rnn_desc

        c_x_descs = _make_tensor_descriptor_array(x_list)
        hx_desc = cudnn.create_tensor_nd_descriptor(hx)
        cx_desc = cudnn.create_tensor_nd_descriptor(cx)

        weights_size = libcudnn.getRNNParamsSize(
{'fraction_non_alphanumeric': 0.10166875663161956, 'fraction_numerical': 0.012539789717372431, 'mean_word_length': 1.9673726388093875, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, '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': '516973', 'n_tokens_mistral': 3629, 'n_tokens_neox': 3429, 'n_words': 746}