<commit_before>"""
Layer definitions
"""

from tf_initializers import ortho_weight, norm_weight
import tensorflow as tf
import numpy 

def matmul3d(x3d, matrix):
    shape = tf.shape(x3d)
    mat_shape = tf.shape(matrix)
    x2d = tf.reshape(x3d, [shape[0]*shape[1], shape[2]])
    result2d = tf.matmul(x2d, matrix)
    result3d = tf.reshape(result2d, [shape[0], shape[1], mat_shape[1]])
    return result3d

class FeedForwardLayer(object):
    def __init__(self,
                 in_size,
                 out_size,
                 non_linearity=tf.nn.tanh):
        self.W = tf.Variable(norm_weight(in_size, out_size), name='W')
        self.b = tf.Variable(numpy.zeros((out_size,)).astype('float32'), name='b')
        self.non_linearity = non_linearity

    def forward(self, x, input_is_3d=False):
        if input_is_3d:
            y = matmul3d(x, self.W) + self.b
        else:
            y = tf.matmul(x, self.W) + self.b
        y = self.non_linearity(y)
        return y

class EmbeddingLayer(object):
    def __init__(self,
                 vocabulary_size,
                 embedding_size):
        self.embeddings = tf.Variable(norm_weight(vocabulary_size, embedding_size),
                                      name='embeddings')
    
    def forward(self, x):
        embs = tf.nn.embedding_lookup(self.embeddings, x)
        return embs

class RecurrentLayer(object):
    def __init__(self,
                 initial_state,
                 step_fn):
        self.initial_state = initial_state
        self.step_fn = step_fn

    def forward(self, x):
        # Assumes that x has shape: time, batch, ...
        states = tf.scan(fn=self.step_fn,
                         elems=x,
                         initializer=self.initial_state)
        return states
    
class GRUStep(object):
    def __init__(self, 
                 input_size, 
                 state_size,
                 input_comes_premultiplied=False):
        self.input_comes_premultiplied = input_comes_premultiplied #TODO
        self.state_to_gates = tf.Variable(
                                numpy.concatenate(
                                    [ortho_weight(state_size),
                                     ortho_weight(state_size)],
                                    axis=1), 
                                name='state_to_gates')
        self.input_to_gates = tf.Variable(
                                numpy.concatenate(
                                    [norm_weight(input_size, state_size),
                                     norm_weight(input_size, state_size)],
                                    axis=1), 
                                name='input_to_gates')
        self.gates_bias = tf.Variable(
                            numpy.zeros((2*state_size,)).astype('float32'),
                            name='gates_bias')

        self.state_to_proposal = tf.Variable(
                                    ortho_weight(state_size),
                                    name = 'state_to_proposal')
        self.input_to_proposal = tf.Variable(
                                    norm_weight(input_size, state_size),
                                    name = 'input_to_proposal')
        self.proposal_bias = tf.Variable(
                                    numpy.zeros((state_size,)).astype('float32'),
                                    name='proposal_bias')

    def _forward(self, prev_state, gates_x, proposal_x):
        gates = tf.matmul(prev_state, self.state_to_gates)
        gates += gates_x
        gates = tf.nn.sigmoid(gates)
        read_gate, update_gate = tf.split(gates,
                                          num_or_size_splits=2,
                                          axis=1)

        proposal = tf.matmul(prev_state, self.state_to_proposal)
        proposal *= read_gate
        proposal += proposal_x
        proposal  = tf.tanh(proposal)

        new_state = update_gate*prev_state + (1-update_gate)*proposal

        return new_state

    def forward(self, prev_state, x):
        #TODO: x*self.input_to_gates can be done in one matmul
        # for all x; same for x*self.input_to_proposal
        gates_x = tf.matmul(x, self.input_to_gates) + self.gates_bias
        proposal_x = tf.matmul(x, self.input_to_proposal) + self.proposal_bias

        return self._forward(prev_state, gates_x, proposal_x)

    def _get_gates_x_proposal_x(self, x):
        gates_x = matmul3d(x, self.input_to_gates) + self.gates_bias
        proposal_x = matmul3d(x, self.input_to_proposal) + self.proposal_bias
        return gates_x, proposal_x


class AttentionStep(object):
    def __init__(self,
                 context,
                 context_state_size,
                 context_mask,
                 state_size,
                 hidden_size):
        self.state_to_hidden = tf.Variable(
                                norm_weight(state_size, hidden_size),
                                name='state_to_hidden')
        self.context_to_hidden = tf.Variable(
                                norm_weight(context_state_size, hidden_size),
                                name='context_to_hidden')
        self.hidden_bias = tf.Variable(
                            numpy.zeros((hidden_size,)).astype('float32'),
                            name='hidden_bias')
        self.hidden_to_score = tf.Variable(
                                norm_weight(hidden_size, 1),
                                name='hidden_to_score')
        
        self.context = context
        self.context_mask = context_mask

        # precompute these activations, they are the same at each step
        # Ideally the compiler would have figured out that too
        self.hidden_from_context = matmul3d(context, self.context_to_hidden)
        self.hidden_from_context += self.hidden_bias

    def forward(self, prev_state):
        hidden = self.hidden_from_context
        hidden += tf.matmul(prev_state, self.state_to_hidden)
        hidden = tf.nn.tanh(hidden)
        # context has shape seqLen x batch x context_state_size
        # mask has shape seqLen x batch

        scores = matmul3d(hidden, self.hidden_to_score) # seqLen x batch x 1
        scores = tf.squeeze(scores, axis=2)
        scores = scores - tf.reduce_max(scores, axis=0, keep_dims=True)
        scores = tf.exp(scores)
        scores *= self.context_mask
        scores = scores / tf.reduce_sum(scores, axis=0, keep_dims=True)

        attention_context = self.context * tf.expand_dims(scores, axis=2)
        attention_context = tf.reduce_sum(attention_context, axis=0, keep_dims=False)

        return attention_context

class Masked_cross_entropy_loss(object):
    def __init__(self,
                 y_true,
                 y_mask):
        self.y_true = tf.reshape(y_true, shape=[-1])
        self.y_mask = y_mask


    def forward(self, logits):
        # TODO: maybe first select the important logits and then compute cost? -- efficiency
        log_probs = tf.nn.log_softmax(logits, dim=-1)
        log_probs = tf.reshape(log_probs, shape=[-1])
        num_classes = tf.shape(logits)[-1] # assumes shape ... x num_classes
        num_logits = tf.size(logits)
        indices = tf.range(start=0, limit=num_logits, delta=num_classes) # limit is excluded
        indices = indices + self.y_true
        log_probs = tf.gather(log_probs, indices)
        log_probs = tf.reshape(log_probs, shape=tf.shape(self.y_mask))
        log_probs *= self.y_mask

        cost = tf.reduce_sum(log_probs, axis=0, keep_dims=False)
        cost = -cost # shape is [batch]

        return cost




        

        


        

                                    

        

                                




<commit_msg>use sparse_softmax_cross_entropy_with_logits - it does the same thing<commit_after>"""
Layer definitions
"""

from tf_initializers import ortho_weight, norm_weight
import tensorflow as tf
import numpy 

def matmul3d(x3d, matrix):
    shape = tf.shape(x3d)
    mat_shape = tf.shape(matrix)
    x2d = tf.reshape(x3d, [shape[0]*shape[1], shape[2]])
    result2d = tf.matmul(x2d, matrix)
    result3d = tf.reshape(result2d, [shape[0], shape[1], mat_shape[1]])
    return result3d

class FeedForwardLayer(object):
    def __init__(self,
                 in_size,
                 out_size,
                 non_linearity=tf.nn.tanh):
        self.W = tf.Variable(norm_weight(in_size, out_size), name='W')
        self.b = tf.Variable(numpy.zeros((out_size,)).astype('float32'), name='b')
        self.non_linearity = non_linearity

    def forward(self, x, input_is_3d=False):
        if input_is_3d:
            y = matmul3d(x, self.W) + self.b
        else:
            y = tf.matmul(x, self.W) + self.b
        y = self.non_linearity(y)
        return y

class EmbeddingLayer(object):
    def __init__(self,
                 vocabulary_size,
                 embedding_size):
        self.embeddings = tf.Variable(norm_weight(vocabulary_size, embedding_size),
                                      name='embeddings')
    
    def forward(self, x):
        embs = tf.nn.embedding_lookup(self.embeddings, x)
        return embs

class RecurrentLayer(object):
    def __init__(self,
                 initial_state,
                 step_fn):
        self.initial_state = initial_state
        self.step_fn = step_fn

    def forward(self, x):
        # Assumes that x has shape: time, batch, ...
        states = tf.scan(fn=self.step_fn,
                         elems=x,
                         initializer=self.initial_state)
        return states
    
class GRUStep(object):
    def __init__(self, 
                 input_size, 
                 state_size,
                 input_comes_premultiplied=False):
        self.input_comes_premultiplied = input_comes_premultiplied #TODO
        self.state_to_gates = tf.Variable(
                                numpy.concatenate(
                                    [ortho_weight(state_size),
                                     ortho_weight(state_size)],
                                    axis=1), 
                                name='state_to_gates')
        self.input_to_gates = tf.Variable(
                                numpy.concatenate(
                                    [norm_weight(input_size, state_size),
                                     norm_weight(input_size, state_size)],
                                    axis=1), 
                                name='input_to_gates')
        self.gates_bias = tf.Variable(
                            numpy.zeros((2*state_size,)).astype('float32'),
                            name='gates_bias')

        self.state_to_proposal = tf.Variable(
                                    ortho_weight(state_size),
                                    name = 'state_to_proposal')
        self.input_to_proposal = tf.Variable(
                                    norm_weight(input_size, state_size),
                                    name = 'input_to_proposal')
        self.proposal_bias = tf.Variable(
                                    numpy.zeros((state_size,)).astype('float32'),
                                    name='proposal_bias')

    def _forward(self, prev_state, gates_x, proposal_x):
        gates = tf.matmul(prev_state, self.state_to_gates)
        gates += gates_x
        gates = tf.nn.sigmoid(gates)
        read_gate, update_gate = tf.split(gates,
                                          num_or_size_splits=2,
                                          axis=1)

        proposal = tf.matmul(prev_state, self.state_to_proposal)
        proposal *= read_gate
        proposal += proposal_x
        proposal  = tf.tanh(proposal)

        new_state = update_gate*prev_state + (1-update_gate)*proposal

        return new_state

    def forward(self, prev_state, x):
        #TODO: x*self.input_to_gates can be done in one matmul
        # for all x; same for x*self.input_to_proposal
        gates_x = tf.matmul(x, self.input_to_gates) + self.gates_bias
        proposal_x = tf.matmul(x, self.input_to_proposal) + self.proposal_bias

        return self._forward(prev_state, gates_x, proposal_x)

    def _get_gates_x_proposal_x(self, x):
        gates_x = matmul3d(x, self.input_to_gates) + self.gates_bias
        proposal_x = matmul3d(x, self.input_to_proposal) + self.proposal_bias
        return gates_x, proposal_x


class AttentionStep(object):
    def __init__(self,
                 context,
                 context_state_size,
                 context_mask,
                 state_size,
                 hidden_size):
        self.state_to_hidden = tf.Variable(
                                norm_weight(state_size, hidden_size),
                                name='state_to_hidden')
        self.context_to_hidden = tf.Variable(
                                norm_weight(context_state_size, hidden_size),
                                name='context_to_hidden')
        self.hidden_bias = tf.Variable(
                            numpy.zeros((hidden_size,)).astype('float32'),
                            name='hidden_bias')
        self.hidden_to_score = tf.Variable(
                                norm_weight(hidden_size, 1),
                                name='hidden_to_score')
        
        self.context = context
        self.context_mask = context_mask

        # precompute these activations, they are the same at each step
        # Ideally the compiler would have figured out that too
        self.hidden_from_context = matmul3d(context, self.context_to_hidden)
        self.hidden_from_context += self.hidden_bias

    def forward(self, prev_state):
        hidden = self.hidden_from_context
        hidden += tf.matmul(prev_state, self.state_to_hidden)
        hidden = tf.nn.tanh(hidden)
        # context has shape seqLen x batch x context_state_size
        # mask has shape seqLen x batch

        scores = matmul3d(hidden, self.hidden_to_score) # seqLen x batch x 1
        scores = tf.squeeze(scores, axis=2)
        scores = scores - tf.reduce_max(scores, axis=0, keep_dims=True)
        scores = tf.exp(scores)
        scores *= self.context_mask
        scores = scores / tf.reduce_sum(scores, axis=0, keep_dims=True)

        attention_context = self.context * tf.expand_dims(scores, axis=2)
        attention_context = tf.reduce_sum(attention_context, axis=0, keep_dims=False)

        return attention_context

class Masked_cross_entropy_loss(object):
    def __init__(self,
                 y_true,
                 y_mask):
        self.y_true = y_true
        self.y_mask = y_mask


    def forward(self, logits):
        # TODO: maybe first select the important logits and then compute cost? -- efficiency

        cost = tf.nn.sparse_softmax_cross_entropy_with_logits(
                labels=self.y_true,
                logits=logits)
        #cost has shape seqLen x batch
        cost *= self.y_mask
        cost = tf.reduce_sum(cost, axis=0, keep_dims=False)
        return cost




        

        


        

                                    

        

                                





{'fraction_non_alphanumeric': 0.08601722995361166, 'fraction_numerical': 0.006295559973492379, 'mean_word_length': 1.6521968365553603, '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': 40, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2807913', 'n_tokens_mistral': 4490, 'n_tokens_neox': 4369, 'n_words': 816}