Add more operation and some refactoring of the code to ease usage.

user_ops/ops/knn_bruteforce.cc

@@ -0,0 +1,41 @@
+// ComputerGraphics Tuebingen, 2018
+
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/shape_inference.h"
+
+namespace tensorflow {
+
+using ::tensorflow::shape_inference::InferenceContext;
+using ::tensorflow::shape_inference::ShapeHandle;
+
+REGISTER_OP("KnnBruteforce")
+    .Input("position: T")  // position:	 each datapoint in nd space [B,Dp, N].
+    .Output("neighborhood_out: NBtype")  // neighborhood_out: all K nearest
+                                         // neighbors        [B, N, K].
+    .Output("distances: T")              // distances: 	 all K nearest distances
+                                         // [B, N, K].
+    .Output("timings: T")                // timings:
+                                         // [1]
+    .Attr("K: int")
+    .Attr("return_timings: bool = false")
+    .Attr("T: realnumbertype")
+    .Attr("NBtype: {int32} = DT_INT32")
+    .SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
+      const auto position = c->input(0);
+      const auto neighborhood_out = c->input(1);
+
+      int K;
+      c->GetAttr("K", &K);
+
+      auto B = c->Dim(position, 0);
+      auto N = c->Dim(position, 2);
+
+      c->set_output(0, c->MakeShape({B, N, K}));
+      c->set_output(1, c->MakeShape({B, N, K}));
+
+      return Status::OK();
+    });
+
+}  // namespace tensorflow
+
+// doc: K:					number of neighbors.