fix shape issues, improve performance, make it pip installable, etc

Author: carusyte

.gitignore

@@ -96,4 +96,8 @@ ENV/
 /site
 
 # mypy
-.mypy_cache/
+.mypy_cache/
+
+# vscode and its extensions
+.vscode/*
+.history/*

dnc/__init__.py

@@ -22,8 +22,8 @@
 import sonnet as snt
 import tensorflow as tf
 
-import addressing
-import util
+from dnc import addressing
+from dnc import util
 
 AccessState = collections.namedtuple('AccessState', (
     'memory', 'read_weights', 'write_weights', 'linkage', 'usage'))
@@ -53,7 +53,7 @@ def _erase_and_write(memory, address, reset_weights, values):
     expand_address = tf.expand_dims(address, 3)
     reset_weights = tf.expand_dims(reset_weights, 2)
     weighted_resets = expand_address * reset_weights
-    reset_gate = tf.reduce_prod(1 - weighted_resets, [1])
+    reset_gate = util.reduce_prod(1 - weighted_resets, 1)
     memory *= reset_gate
 
   with tf.name_scope('additive_write', values=[memory, address, values]):

access.py dnc/access.pyaccess.py renamed to dnc/access.py

@@ -22,8 +22,8 @@
 import tensorflow as tf
 from tensorflow.python.ops import rnn
 
-import access
-import util
+from dnc import access
+from dnc import util
 
 BATCH_SIZE = 2
 MEMORY_SIZE = 20

access_test.py dnc/access_test.pyaccess_test.py renamed to dnc/access_test.py

@@ -22,7 +22,7 @@
 import sonnet as snt
 import tensorflow as tf
 
-import util
+from dnc import util
 
 # Ensure values are greater than epsilon to avoid numerical instability.
 _EPSILON = 1e-6
@@ -32,7 +32,7 @@
 
 
 def _vector_norms(m):
-  squared_norms = tf.reduce_sum(m * m, axis=2, keep_dims=True)
+  squared_norms = tf.reduce_sum(m * m, axis=2, keepdims=True)
   return tf.sqrt(squared_norms + _EPSILON)
 
 
@@ -202,7 +202,7 @@ def _link(self, prev_link, prev_precedence_weights, write_weights):
       containing the new link graphs for each write head.
     """
     with tf.name_scope('link'):
-      batch_size = prev_link.get_shape()[0].value
+      batch_size = tf.shape(prev_link)[0]
       write_weights_i = tf.expand_dims(write_weights, 3)
       write_weights_j = tf.expand_dims(write_weights, 2)
       prev_precedence_weights_j = tf.expand_dims(prev_precedence_weights, 2)
@@ -236,7 +236,7 @@ def _precedence_weights(self, prev_precedence_weights, write_weights):
       new precedence weights.
     """
     with tf.name_scope('precedence_weights'):
-      write_sum = tf.reduce_sum(write_weights, 2, keep_dims=True)
+      write_sum = tf.reduce_sum(write_weights, 2, keepdims=True)
       return (1 - write_sum) * prev_precedence_weights + write_weights
 
   @property
@@ -351,7 +351,7 @@ def _usage_after_write(self, prev_usage, write_weights):
     """
     with tf.name_scope('usage_after_write'):
       # Calculate the aggregated effect of all write heads
-      write_weights = 1 - tf.reduce_prod(1 - write_weights, [1])
+      write_weights = 1 - util.reduce_prod(1 - write_weights, 1)
       return prev_usage + (1 - prev_usage) * write_weights
 
   def _usage_after_read(self, prev_usage, free_gate, read_weights):
@@ -370,7 +370,7 @@ def _usage_after_read(self, prev_usage, free_gate, read_weights):
     with tf.name_scope('usage_after_read'):
       free_gate = tf.expand_dims(free_gate, -1)
       free_read_weights = free_gate * read_weights
-      phi = tf.reduce_prod(1 - free_read_weights, [1], name='phi')
+      phi = util.reduce_prod(1 - free_read_weights, 1, name='phi')
       return prev_usage * phi
 
   def _allocation(self, usage):

addressing.py dnc/addressing.pyaddressing.py renamed to dnc/addressing.py

@@ -22,8 +22,8 @@
 import sonnet as snt
 import tensorflow as tf
 
-import addressing
-import util
+from dnc import addressing
+from dnc import util
 
 
 class WeightedSoftmaxTest(tf.test.TestCase):

addressing_test.py dnc/addressing_test.pyaddressing_test.py renamed to dnc/addressing_test.py

@@ -27,7 +27,7 @@
 import sonnet as snt
 import tensorflow as tf
 
-import access
+from dnc import access
 
 DNCState = collections.namedtuple('DNCState', ('access_output', 'access_state',
                                                'controller_state'))
@@ -110,7 +110,7 @@ def _build(self, inputs, prev_state):
         controller_input, prev_controller_state)
 
     controller_output = self._clip_if_enabled(controller_output)
-    controller_state = snt.nest.map(self._clip_if_enabled, controller_state)
+    controller_state = tf.contrib.framework.nest.map_structure(self._clip_if_enabled, controller_state)
 
     access_output, access_state = self._access(controller_output,
                                                prev_access_state)

dnc.py dnc/dnc.pydnc.py renamed to dnc/dnc.py

@@ -25,21 +25,48 @@
 def batch_invert_permutation(permutations):
   """Returns batched `tf.invert_permutation` for every row in `permutations`."""
   with tf.name_scope('batch_invert_permutation', values=[permutations]):
-    unpacked = tf.unstack(permutations)
-    inverses = [tf.invert_permutation(permutation) for permutation in unpacked]
-    return tf.stack(inverses)
+    perm = tf.cast(permutations, tf.float32)
+    dim = int(perm.get_shape()[-1])
+    size = tf.cast(tf.shape(perm)[0], tf.float32)
+    delta = tf.cast(tf.shape(perm)[-1], tf.float32)
+    rg = tf.range(0, size * delta, delta, dtype=tf.float32)
+    rg = tf.expand_dims(rg, 1)
+    rg = tf.tile(rg, [1, dim])
+    perm = tf.add(perm, rg)
+    flat = tf.reshape(perm, [-1])
+    perm = tf.invert_permutation(tf.cast(flat, tf.int32))
+    perm = tf.reshape(perm, [-1, dim])
+    return tf.subtract(perm, tf.cast(rg, tf.int32))
 
 
 def batch_gather(values, indices):
   """Returns batched `tf.gather` for every row in the input."""
   with tf.name_scope('batch_gather', values=[values, indices]):
-    unpacked = zip(tf.unstack(values), tf.unstack(indices))
-    result = [tf.gather(value, index) for value, index in unpacked]
-    return tf.stack(result)
+    idx = tf.expand_dims(indices, -1)
+    size = tf.shape(indices)[0]
+    rg = tf.range(size, dtype=tf.int32)
+    rg = tf.expand_dims(rg, -1)
+    rg = tf.tile(rg, [1, int(indices.get_shape()[-1])])
+    rg = tf.expand_dims(rg, -1)
+    gidx = tf.concat([rg, idx], -1)
+    return tf.gather_nd(values, gidx)
 
 
 def one_hot(length, index):
   """Return an nd array of given `length` filled with 0s and a 1 at `index`."""
   result = np.zeros(length)
   result[index] = 1
   return result
+
+def reduce_prod(x, axis, name=None):
+  """Efficient reduce product over axis.
+
+  Uses tf.cumprod and tf.gather_nd as a workaround to the poor performance of calculating tf.reduce_prod's gradient on CPU.
+  """
+  with tf.name_scope(name, 'util_reduce_prod', values=[x]):
+    cp = tf.cumprod(x, axis, reverse=True)
+    size = tf.shape(cp)[0]
+    idx1 = tf.range(tf.cast(size, tf.float32), dtype=tf.float32)
+    idx2 = tf.zeros([size], tf.float32)
+    indices = tf.stack([idx1, idx2], 1)
+    return tf.gather_nd(cp, tf.cast(indices, tf.int32))