<commit_before>import tensorflow as tf
import numpy as np
import unittest

from dnc.memory import Memory

class DNCMemoryTests(unittest.TestCase):

    def test_construction(self):
        graph = tf.Graph()
        with graph.as_default():
            with tf.Session(graph=graph) as session:

                mem = Memory(4, 5, 2)
                session.run(tf.initialize_all_variables())

                self.assertEqual(mem.words_num, 4)
                self.assertEqual(mem.word_size, 5)
                self.assertEqual(mem.read_heads, 2)

                self.assertTrue(isinstance(mem.memory_matrix, tf.Variable))
                self.assertEqual(mem.memory_matrix.get_shape().as_list(), [4, 5])
                self.assertTrue(isinstance(mem.usage_vector, tf.Variable))
                self.assertEqual(mem.usage_vector.get_shape().as_list(), [4])
                self.assertTrue(isinstance(mem.link_matrix, tf.Variable))
                self.assertEqual(mem.link_matrix.get_shape().as_list(), [4, 4])
                self.assertTrue(isinstance(mem.write_weighting, tf.Variable))
                self.assertEqual(mem.write_weighting.get_shape().as_list(), [4])
                self.assertTrue(isinstance(mem.read_weightings, tf.Variable))
                self.assertEqual(mem.read_weightings.get_shape().as_list(), [2, 4])
                self.assertTrue(isinstance(mem.read_vectors, tf.Variable))
                self.assertEqual(mem.read_vectors.get_shape().as_list(), [2, 5])


    def test_lookup_weighting(self):
        graph = tf.Graph()
        with graph.as_default():
            with tf.Session(graph=graph) as session:

                mem = Memory(4, 5, 2)
                keys = np.array([[0., 1., 0., 0.3, 4.3], [1.3, 0.8, 0., 0., 0.62]]).astype(np.float32)
                strengths = np.array([0.7, 0.2]).astype(np.float32)
                predicted_weights = np.array([[0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25]]).astype(np.float32)

                op = mem.get_lookup_weighting(keys, strengths)
                session.run(tf.initialize_all_variables())
                c = session.run(op)

                self.assertTrue(c.shape, (2, 4))
                self.assertTrue(np.array_equal(c, predicted_weights))

if __name__ == '__main__':
    unittest.main(verbosity=2)<commit_msg>added test for update_usage_vector<commit_after>import tensorflow as tf
import numpy as np
import unittest

from dnc.memory import Memory

class DNCMemoryTests(unittest.TestCase):

    def test_construction(self):
        graph = tf.Graph()
        with graph.as_default():
            with tf.Session(graph=graph) as session:

                mem = Memory(4, 5, 2)
                session.run(tf.initialize_all_variables())

                self.assertEqual(mem.words_num, 4)
                self.assertEqual(mem.word_size, 5)
                self.assertEqual(mem.read_heads, 2)

                self.assertTrue(isinstance(mem.memory_matrix, tf.Variable))
                self.assertEqual(mem.memory_matrix.get_shape().as_list(), [4, 5])
                self.assertTrue(isinstance(mem.usage_vector, tf.Variable))
                self.assertEqual(mem.usage_vector.get_shape().as_list(), [4])
                self.assertTrue(isinstance(mem.link_matrix, tf.Variable))
                self.assertEqual(mem.link_matrix.get_shape().as_list(), [4, 4])
                self.assertTrue(isinstance(mem.write_weighting, tf.Variable))
                self.assertEqual(mem.write_weighting.get_shape().as_list(), [4])
                self.assertTrue(isinstance(mem.read_weightings, tf.Variable))
                self.assertEqual(mem.read_weightings.get_shape().as_list(), [2, 4])
                self.assertTrue(isinstance(mem.read_vectors, tf.Variable))
                self.assertEqual(mem.read_vectors.get_shape().as_list(), [2, 5])


    def test_lookup_weighting(self):
        graph = tf.Graph()
        with graph.as_default():
            with tf.Session(graph=graph) as session:

                mem = Memory(4, 5, 2)
                keys = np.array([[0., 1., 0., 0.3, 4.3], [1.3, 0.8, 0., 0., 0.62]]).astype(np.float32)
                strengths = np.array([0.7, 0.2]).astype(np.float32)
                predicted_weights = np.array([[0.25, 0.25, 0.25, 0.25], [0.25, 0.25, 0.25, 0.25]]).astype(np.float32)

                op = mem.get_lookup_weighting(keys, strengths)
                session.run(tf.initialize_all_variables())
                c = session.run(op)

                self.assertTrue(c.shape, (2, 4))
                self.assertTrue(np.array_equal(c, predicted_weights))


    def test_update_usage_vector(self):
        graph = tf.Graph()
        with graph.as_default():
            with tf.Session(graph=graph) as session:

                mem = Memory(4, 5, 2)
                free_gates = np.array([0.2, 0.67]).astype(np.float32)
                predicted_usage = np.array([0.49429685,  0.49429685,  0.49429685,  0.49429685]).astype(np.float32)

                changes = [
                    mem.read_weightings.assign(tf.fill([2, 4], 0.25)),
                    mem.write_weighting.assign(tf.fill([4, ], 0.25)),
                    mem.usage_vector.assign(tf.fill([4, ], 0.5))
                ]
                op = mem.update_usage_vector(free_gates)
                session.run(tf.initialize_all_variables())
                session.run(changes)
                u = session.run(op)
                updated_usage = session.run(mem.usage_vector.value())

                self.assertEqual(u.shape, (4, ))
                self.assertTrue(np.array_equal(u, predicted_usage))
                self.assertTrue(np.array_equal(updated_usage, predicted_usage))

if __name__ == '__main__':
    unittest.main(verbosity=2)
{'fraction_non_alphanumeric': 0.133785664578984, 'fraction_numerical': 0.03566457898399443, 'mean_word_length': 2.1143011917659806, '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': '5246820', 'n_tokens_mistral': 1924, 'n_tokens_neox': 1744, 'n_words': 280}