Added inhibit C++ op

.gitignore

@@ -1,2 +1,3 @@
 __pycache__
 data
+*.so

compile.sh

@@ -0,0 +1,3 @@
+#!/bin/bash
+TF_INC=$(python3.5 -c 'import tensorflow as tf; print(tf.sysconfig.get_include())')
+g++ --std=c++11 --shared inhibit.cc -o inhibit.so -fPIC -I $TF_INC -O2

inhibit.cc

@@ -0,0 +1,78 @@
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/shape_inference.h"
+#include "tensorflow/core/framework/op_kernel.h"
+
+#include <tuple>
+#include <queue>
+#include <unordered_set>
+#include <math.h>
+
+using namespace std;
+using namespace tensorflow;
+
+REGISTER_OP("Inhibit")
+    .Input("to_inhibit: float")
+    .Output("inhibited: float")
+    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
+      c->set_output(0, c->input(0));
+      return Status::OK();
+    });
+
+/**
+ * Applies inhibition to a matrix, setting all but the highest values
+ * in the matrix to zero, and all highest values to one.
+ */
+class InhibitOp : public OpKernel {
+ public:
+  explicit InhibitOp(OpKernelConstruction* context) : OpKernel(context) {}
+
+  void Compute(OpKernelContext* context) override {
+    // Grab the input tensor
+    const Tensor& input_tensor = context->input(0);
+    auto input = input_tensor.flat<float>();
+
+    // Create an output tensor
+    Tensor* output_tensor = NULL;
+    OP_REQUIRES_OK(context, context->allocate_output(0, input_tensor.shape(),
+                                                     &output_tensor));
+    auto output = output_tensor->flat<float>();
+
+    // Set all but the first element of the output tensor to 0.
+    const int N = input.size();
+    const int maxActive = (int) ceil(N * 0.2f);
+
+    priority_queue<tuple<int, int>> keep;
+
+    for (int i = 0; i < N; i++) {
+      auto score = input(i);
+
+      if(keep.empty() || score > get<0>(keep.top())) {
+        // We want to keep this
+        if(keep.size() >= maxActive) {
+          // Replace lowest score with this one
+          keep.pop();
+        }
+
+        keep.push(make_tuple(score, i));
+      }
+    }
+
+    // Create a set of indicies to keep
+    unordered_set<int> keepIndicies;
+
+    while(!keep.empty()) {
+      keepIndicies.insert(get<1>(keep.top()));
+      keep.pop();
+    }
+
+    // Handle output
+    for (int i = 0; i < N; i++) {
+      if(keepIndicies.find(i) != keepIndicies.end())
+        output(i) = 1;
+      else
+        output(i) = 0;
+    }
+  }
+};
+
+REGISTER_KERNEL_BUILDER(Name("Inhibit").Device(DEVICE_CPU), InhibitOp);

main2.py

@@ -3,6 +3,8 @@
 from pq import PriorityQueue
 from math import *
 
+inhibit_module = tf.load_op_library('./inhibit.so')
+
 epochs = 1
 dim = (10, 10)
 sparsity = 0.2
@@ -31,31 +33,6 @@ def tf_count(t, val):
     count = tf.reduce_sum(as_ints)
     return count
 
-def global_inhibit(overlap_scores, out_size, max_active):
-    # Set of outputs to keep
-    keep = PriorityQueue()
-
-    for i in range(out_size):
-        score = overlap_scores[0][i]
-
-        if keep.empty() or score > keep.peek():
-            # We want this score to be kept
-
-            if keep.size() < max_active:
-                keep.update(i, score)
-            else:
-                # Replace the lowest score in keep with this one
-                keep.pushpop(i, score)
-
-    # Set all indices not kept to 0
-    keep_indicies = {i for _, _, i in keep.heap}
-
-    for i in range(out_size):
-        if i not in keep_indicies:
-            overlap_scores[0][i] = 0
-
-    return overlap_scores
-
 # Defines the condition to conert permanence into connections
 p_cond = tf.greater(p, 0)
 
@@ -65,8 +42,8 @@ def global_inhibit(overlap_scores, out_size, max_active):
 # Compute the activation before inhibition applied
 activation = tf.matmul(x, c)
 
-# TODO: Output vector
-y = tf.while_loop(lambda a: tf_count(a, 0) < output_zeros, lambda a: tf.sub(a, tf.ones(1)), [activation])
+# Output vector
+y = inhibit_module.inhibit(activation)
 
 init_op = tf.initialize_all_variables()
 
@@ -87,5 +64,5 @@ def global_inhibit(overlap_scores, out_size, max_active):
     sess.run(update_connections)
 
     for _ in range(1):
-        result = sess.run([activation], feed_dict={x: [[0,0,0,0,1,0,0,0,1,0]]})
-        print(global_inhibit(result, 10, 2))
+        result, = sess.run([y], feed_dict={x: [[0,0,0,0,1,0,0,0,1,0]]})
+        print(result)