From 2cbeb7e69953fb4ef42c7bb53e7d93262f4b9561 Mon Sep 17 00:00:00 2001
From: Fabian Groh <fabian.groh@gmail.com>
Date: Wed, 31 Oct 2018 14:11:10 +0100
Subject: [PATCH] Add more operation and some refactoring of the code to ease
 usage.

---
 example.py                                    |  40 ++-
 layers.py                                     | 100 ++++++-
 user_ops/CMakeLists.txt                       |   8 +
 user_ops/FindTensorFlow.cmake                 |  21 +-
 user_ops/PointTestCase.py                     | 130 ---------
 user_ops/__init__.py                          |  64 ++--
 user_ops/kernels/cuda_utils.h                 |  10 +
 user_ops/kernels/flex_conv_kernel.cc          |   3 +-
 user_ops/kernels/flex_conv_kernel_gpu.cu.cc   |  37 ++-
 user_ops/kernels/flex_conv_op.cc              |  42 +--
 user_ops/kernels/flex_deconv_kernel.cc        |   4 +-
 user_ops/kernels/flex_deconv_kernel_gpu.cu.cc |   2 +
 user_ops/kernels/flex_deconv_op.cc            |  28 +-
 user_ops/kernels/flex_pool_kernel.cc          |   6 +-
 user_ops/kernels/flex_pool_kernel_gpu.cu.cc   |  30 +-
 user_ops/kernels/flex_pool_op.cc              |  40 +--
 user_ops/kernels/knn_bruteforce_kernel.cc     |  79 +++++
 .../kernels/knn_bruteforce_kernel_gpu.cu.cc   | 273 ++++++++++++++++++
 user_ops/kernels/knn_bruteforce_op.cc         |  81 ++++++
 user_ops/kernels/knn_bruteforce_op.h          |  44 +++
 user_ops/misc.py                              | 128 ++++++++
 user_ops/ops/knn_bruteforce.cc                |  41 +++
 user_ops/test_all.py                          |  29 ++
 user_ops/test_flex_convolution.py             | 150 ++++++----
 user_ops/test_flex_convolution_transpose.py   | 151 +++++-----
 user_ops/test_flex_pooling.py                 | 109 ++++---
 user_ops/test_knn_bruteforce.py               |  66 +++++
 27 files changed, 1264 insertions(+), 452 deletions(-)
 delete mode 100644 user_ops/PointTestCase.py
 create mode 100644 user_ops/kernels/cuda_utils.h
 create mode 100644 user_ops/kernels/knn_bruteforce_kernel.cc
 create mode 100644 user_ops/kernels/knn_bruteforce_kernel_gpu.cu.cc
 create mode 100644 user_ops/kernels/knn_bruteforce_op.cc
 create mode 100644 user_ops/kernels/knn_bruteforce_op.h
 create mode 100644 user_ops/misc.py
 create mode 100644 user_ops/ops/knn_bruteforce.cc
 create mode 100644 user_ops/test_all.py
 create mode 100644 user_ops/test_knn_bruteforce.py

diff --git a/example.py b/example.py
index 11b152c..2b44cf5 100644
--- a/example.py
+++ b/example.py
@@ -25,39 +25,51 @@ Demonstration of using FlexConvolution, FlexPooling Layer.
 import numpy as np
 import tensorflow as tf
 from tabulate import tabulate
-from layers import flex_convolution, flex_convolution_transpose, flex_pooling
+from layers import (flex_convolution,
+                    flex_convolution_transpose,
+                    flex_pooling,
+                    knn_bruteforce)
 
-B, Din, Dout, Dout2, Dp, N, N2, K, K2 = 1, 2, 4, 8, 3, 10, 5, 5, 3
+
+B, Din, Dout, Dout2, Dp, N, K = 1, 2, 4, 8, 3, 10, 5
 
 features = np.random.randn(B, Din, N).astype(np.float32)
 positions = np.random.randn(B, Dp, N).astype(np.float32)
-neighbors = np.random.randint(0, N, [B, K, N]).astype(np.int32)
-neighbors2 = np.random.randint(0, N, [B, K2, N2]).astype(np.int32)
 
 features = tf.convert_to_tensor(features, name='features')
 positions = tf.convert_to_tensor(positions, name='positions')
-neighbors = tf.convert_to_tensor(neighbors, name='neighbors')
-neighbors2 = tf.convert_to_tensor(neighbors2, name='neighbors2')
 
 net = [features]
 # use our FlexConv similar to a traditional convolution layer
-net.append(flex_convolution(net[-1], positions, neighbors, Dout,
+
+neighbors = knn_bruteforce(positions, K=5)
+net.append(flex_convolution(net[-1],
+                            positions,
+                            neighbors,
+                            Dout,
                             activation=tf.nn.relu))
 # pool and sub-sampling are different operations
 net.append(flex_pooling(net[-1], neighbors))
 
 # when ordering the points beforehand sub-sampling is simply
-features = net[-1][:, :, :N2]
-positions = positions[:, :, :N2]
-
+features = net[-1][:, :, :N // 2]
+positions = positions[:, :, :N // 2]
 net.append(features)
+
+neighbors = knn_bruteforce(positions, K=3)
 # we didn't notice any improvements using the transposed version vs. pooling
-net.append(flex_convolution_transpose(net[-1], positions, neighbors2, Dout2,
+net.append(flex_convolution_transpose(net[-1],
+                                      positions,
+                                      neighbors,
+                                      Dout,
                                       activation=tf.nn.relu))
 # of course any commonly used arguments work here as well
-net.append(flex_convolution(net[-1], positions,
-                            neighbors2, Dout2,
-                            trainable=False, activation=tf.nn.relu))
+net.append(flex_convolution(net[-1],
+                            positions,
+                            neighbors,
+                            Dout,
+                            activation=tf.nn.relu,
+                            trainable=False))
 
 gradient_wrt_feature = tf.gradients(net[-1], net[0])
 
diff --git a/layers.py b/layers.py
index 93fc582..fe9feaa 100644
--- a/layers.py
+++ b/layers.py
@@ -20,25 +20,86 @@
 import tensorflow as tf
 from user_ops import flex_convolution as _flex_convolution
 from user_ops import flex_pooling as _flex_pooling
+from user_ops import knn_bruteforce as _knn_bruteforce
 from user_ops import flex_convolution_transpose as _flex_convolution_transpose
 
 from tensorflow.python.keras import activations
 from tensorflow.python.keras import initializers
 from tensorflow.python.layers.base import Layer
-from tensorflow.python.util.tf_export import tf_export
 from tensorflow.python.framework import tensor_shape
 from tensorflow.python.framework import ops
 
 
 all = ['FlexPooling', 'FlexConvolution', 'FlexConvolutionTranspose',
-       'flex_pooling', 'flex_convolution', 'flex_convolution_transpose']
+       'flex_pooling', 'flex_convolution', 'flex_convolution_transpose',
+       'knn_bruteforce', 'KnnBruteforce']
 
 
 def _remove_dim(x, axis=2):
   return tf.squeeze(x, axis=axis)
 
 
-@tf_export('keras.layers.FlexPooling')
+class KnnBruteforce(Layer):
+  """knn bruteforce layer.
+  This layer performs a nearest neighbor lookup for a batch of given positions.
+
+  Arguments:
+      K: size of each neighborhood
+      data_format: A string, one of `simple` (default) or `expanded`.
+        If `simple` the shapes are [B, Din, N], when `expanded` the shapes
+        are assumed to be [B, Din, 1, N] to match `channels_first` in trad
+        convolutions.
+      name: A string, the name of the layer.
+
+  Inputs:
+      positions: A `Tensor` of the format [B, Dp, (1), N].
+
+  Outputs:
+      neighborhoods: A `Tensor` of the format [B, K, (1), N] (tf.int32).
+
+  """
+
+  def __init__(self,
+               K,
+               data_format='simple',
+               name=None):
+
+    super(KnnBruteforce, self).__init__(name=name)
+    assert K > 0
+    assert data_format in ['simple', 'expanded']
+    self.K = K
+    self.data_format = data_format
+
+  def compute_output_shape(self, input_shapes):
+    output_shape = input_shapes[0]
+    output_shape[1] = self.K
+    return output_shape
+
+  def call(self, inputs):
+
+    positions = ops.convert_to_tensor(inputs, dtype=self.dtype)
+    if self.data_format == 'expanded':
+      positions = _remove_dim(positions, 2)
+
+    NN, _, _ = _knn_bruteforce(positions, K=4)
+    NN = tf.transpose(NN, [0, 2, 1])
+
+    if self.data_format == 'expanded':
+      NN = tf.expand_dims(NN, axis=2)
+
+    return NN
+
+
+def knn_bruteforce(positions,
+                   K,
+                   data_format='simple',
+                   name=None):
+
+  layer = KnnBruteforce(K, data_format=data_format, name=name)
+
+  return layer.apply(positions)
+
+
 class FlexPooling(Layer):
   """flex pooling layer.
 
@@ -53,8 +114,19 @@ class FlexPooling(Layer):
   Arguments:
       features: A `Tensor` of the format [B, Din, (1), N].
       neighborhoods: A `Tensor` of the format [B, K, (1), N] (tf.int32).
+      data_format: A string, one of `simple` (default) or `expanded`.
+        If `simple` the shapes are [B, Din, N], when `expanded` the shapes
+        are assumed to be [B, Din, 1, N] to match `channels_first` in trad
+        convolutions.
       name: A string, the name of the layer.
 
+  Inputs:
+      features: A `Tensor` of the format[B, Din, (1), N].
+      neighborhoods: A `Tensor` of the format [B, K, (1), N] (tf.int32).
+
+  Outputs:
+      features: A `Tensor` of the format[B, Din, (1), N].
+
   """
 
   def __init__(self,
@@ -62,6 +134,7 @@ class FlexPooling(Layer):
                name=None):
 
     super(FlexPooling, self).__init__(name=name)
+    assert data_format in ['simple', 'expanded']
     self.data_format = data_format
 
   def compute_output_shape(self, input_shape):
@@ -93,11 +166,9 @@ def flex_pooling(features,
                  name=None):
 
   layer = FlexPooling(data_format=data_format, name=name)
-
   return layer.apply([features, neighborhoods])
 
 
-@tf_export('keras.layers.FlexConvolution')
 class FlexConvolution(Layer):
   """flex convolution layer.
 
@@ -126,7 +197,7 @@ class FlexConvolution(Layer):
       features_bias_initializer: An initializer for the bias vector after
         the convolution. If None, the default initializer will be used.
       use_feature_bias: Boolean, whether the layer uses a bias.
-      data_format: A string, one of `simple` (default) or `expaned`.
+      data_format: A string, one of `simple` (default) or `expanded`.
         If `simple` the shapes are [B, Din, N], when `expanded` the shapes
         are assumed to be [B, Din, 1, N] to match `channels_first` in trad
         convolutions.
@@ -136,9 +207,11 @@ class FlexConvolution(Layer):
 
   Inputs:
       features: A `Tensor` of the format[B, Din, (1), N].
-  positions: A `Tensor` of the format[B, Dp, (1), N].
-  neigh
-borhoods: A `Tensor` of the format [B, K, (1), N] (tf.int32).
+      positions: A `Tensor` of the format[B, Dp, (1), N].
+      neighborhoods: A `Tensor` of the format [B, K, (1), N] (tf.int32).
+
+  Outputs:
+      features: A `Tensor` of the format[B, Dout, (1), N].
 
   """
 
@@ -155,6 +228,7 @@ borhoods: A `Tensor` of the format [B, K, (1), N] (tf.int32).
 
     super(FlexConvolution, self).__init__(trainable=trainable,
                                           name=name)
+    assert data_format in ['simple', 'expanded']
     self.filters = int(filters)
     self.activation = activations.get(activation)
     self.use_feature_bias = use_feature_bias
@@ -226,7 +300,7 @@ borhoods: A `Tensor` of the format [B, K, (1), N] (tf.int32).
 
     Returns:
         tf.tensor: A `Tensor` of the format [B, Dout, (1), N] describing
-        the outgoing features.
+          the outgoing features.
 
     """
     if not isinstance(inputs, list):
@@ -283,7 +357,6 @@ def flex_convolution(features,
   return layer.apply([features, positions, neighborhoods])
 
 
-@tf_export('keras.layers.FlexConvolutionTranspose')
 class FlexConvolutionTranspose(FlexConvolution):
   """flex convolution-transpose layer.
 
@@ -312,7 +385,7 @@ class FlexConvolutionTranspose(FlexConvolution):
       features_bias_initializer: An initializer for the bias vector after
         the convolution. If None, the default initializer will be used.
       use_feature_bias: Boolean, whether the layer uses a bias.
-      data_format: A string, one of `simple` (default) or `expaned`.
+      data_format: A string, one of `simple` (default) or `expanded`.
         If `simple` the shapes are [B, Din, N], when `expanded` the shapes
         are assumed to be [B, Din, 1, N] to match `channels_first` in trad
         convolutions.
@@ -325,6 +398,9 @@ class FlexConvolutionTranspose(FlexConvolution):
       positions: A `Tensor` of the format [B, Dp, (1), N].
       neighborhoods: A `Tensor` of the format [B, K, (1), N] (tf.int32).
 
+  Outputs:
+      features: A `Tensor` of the format[B, Dout, (1), N].
+
   """
 
   def internal_call(self,
diff --git a/user_ops/CMakeLists.txt b/user_ops/CMakeLists.txt
index 33e0de0..25da6b5 100644
--- a/user_ops/CMakeLists.txt
+++ b/user_ops/CMakeLists.txt
@@ -13,7 +13,12 @@ list(APPEND CMAKE_MODULE_PATH ${PROJECT_SOURCE_DIR})
 find_package(CUDA REQUIRED)
 find_package(TensorFlow REQUIRED)
 
+# include CUDA Toolkit header files
+set(CUDA_SAMPLE_INC "$ENV{CUDA_INSTALL_PATH}/samples/common/inc")
+message(STATUS "CUDA_INCLUDE_DIRS: ${CUDA_INCLUDE_DIRS}")
+message(STATUS "CUDA_SAMPLE_INC: ${CUDA_SAMPLE_INC}")
 
+# include CUBS
 if (DEFINED ENV{CUB_INC})
     message(STATUS "Use Cuda-CUB from " $ENV{CUB_INC})
     set(CUB_INC $ENV{CUB_INC})
@@ -31,6 +36,8 @@ include_directories(SYSTEM "/usr/local/")
 # fix cgtuebingen
 include_directories(SYSTEM "/graphics/opt/opt_Ubuntu18.04/cuda/toolkit_9.2")
 
+include_directories(SYSTEM "${CUDA_SAMPLE_INC}")
+include_directories(SYSTEM "${CUDA_INCLUDE_DIRS}/../../")
 include_directories(SYSTEM ${CUB_INC})
 include_directories(SYSTEM ${TensorFlow_INCLUDE_DIR})
 include_directories(SYSTEM kernels)
@@ -38,3 +45,4 @@ include_directories(SYSTEM kernels)
 add_tensorflow_gpu_operation("flex_conv")
 add_tensorflow_gpu_operation("flex_deconv")
 add_tensorflow_gpu_operation("flex_pool")
+add_tensorflow_gpu_operation("knn_bruteforce")
diff --git a/user_ops/FindTensorFlow.cmake b/user_ops/FindTensorFlow.cmake
index 41c9b04..58ecda1 100644
--- a/user_ops/FindTensorFlow.cmake
+++ b/user_ops/FindTensorFlow.cmake
@@ -54,6 +54,7 @@ endif(WIN32)
 
 set(PYTHON_EXECUTABLE "python3" CACHE STRING "specify the python version TensorFlow is installed on.")
 
+
 if(TensorFlow_FOUND)
   # reuse cached variables
   message(STATUS "Reuse cached information from TensorFlow ${TensorFlow_VERSION} ")
@@ -122,7 +123,7 @@ else()
     # However, only TensorFlow versions 1.9, 1.10 support all header files
     # for custom ops.
     set(_TensorFlow_KNOWN_VERSIONS ${TensorFlow_ADDITIONAL_VERSIONS}
-        "1.9" "1.9.0" "1.10" "1.10.0", "1.11", "1.11.0")
+        "1.9" "1.9.0" "1.10" "1.10.0" "1.11" "1.11.0")
     set(_TensorFlow_TEST_VERSIONS)
 
     if(TF_FIND_VERSION)
@@ -166,6 +167,17 @@ else()
                       "We tested against ${_TensorFlow_TEST_VERSIONS}")
   endif(NOT TensorFlow_FOUND)
 
+  # test 1.11 version
+  if("${TF_DETECTED_VERSION}" VERSION_EQUAL "1.11")
+    set(TF_DISABLE_ASSERTS "TRUE")
+  endif()
+
+endif()
+
+
+if(${TF_DISABLE_ASSERTS})
+  message(STATUS "[WARNING] The TensorFlow version ${TF_DETECTED_VERSION} has a bug (see \#22766). We disable asserts using -DNDEBUG=True ")
+  add_definitions("-DNDEBUG=True")
 endif()
 
 find_library(TensorFlow_C_LIBRARY
@@ -246,8 +258,8 @@ endmacro()
 
 # simplify TensorFlow dependencies
 add_library(TensorFlow_DEP INTERFACE)
-TARGET_INCLUDE_DIRECTORIES(TensorFlow_DEP INTERFACE ${TensorFlow_SOURCE_DIR})
-TARGET_INCLUDE_DIRECTORIES(TensorFlow_DEP INTERFACE ${TensorFlow_INCLUDE_DIR})
+TARGET_INCLUDE_DIRECTORIES(TensorFlow_DEP SYSTEM INTERFACE ${TensorFlow_SOURCE_DIR})
+TARGET_INCLUDE_DIRECTORIES(TensorFlow_DEP SYSTEM INTERFACE ${TensorFlow_INCLUDE_DIR})
 TARGET_LINK_LIBRARIES(TensorFlow_DEP INTERFACE -Wl,--allow-multiple-definition -Wl,--whole-archive ${TensorFlow_C_LIBRARY} -Wl,--no-whole-archive)
 TARGET_LINK_LIBRARIES(TensorFlow_DEP INTERFACE -Wl,--allow-multiple-definition -Wl,--whole-archive ${TensorFlow_LIBRARY} -Wl,--no-whole-archive)
 
@@ -263,7 +275,7 @@ find_package_handle_standard_args(
   )
 
 mark_as_advanced(TF_INFORMATION_STRING TF_DETECTED_VERSION TF_DETECTED_VERSION_MAJOR TF_DETECTED_VERSION_MINOR TF_DETECTED_VERSION TF_DETECTED_ABI
-                 TF_DETECTED_INCLUDE_DIR TF_DETECTED_LIBRARY
+                 TF_DETECTED_INCLUDE_DIR TF_DETECTED_LIBRARY TF_DISABLE_ASSERTS
                  TensorFlow_C_LIBRARY TensorFlow_LIBRARY TensorFlow_SOURCE_DIR TensorFlow_INCLUDE_DIR TensorFlow_ABI)
 
 SET(TensorFlow_INCLUDE_DIR ${TensorFlow_INCLUDE_DIR} CACHE PATH "path to tensorflow header files")
@@ -271,3 +283,4 @@ SET(TensorFlow_VERSION ${TensorFlow_VERSION} CACHE INTERNAL "The Python executab
 SET(TensorFlow_ABI ${TensorFlow_ABI} CACHE STRING "The Python executable Version")
 SET(TensorFlow_LIBRARY ${TensorFlow_LIBRARY} CACHE PATH "The Python executable Version")
 SET(TensorFlow_FOUND ${TensorFlow_FOUND} CACHE BOOL "The Python executable Version")
+SET(TF_DISABLE_ASSERTS ${TF_DISABLE_ASSERTS} CACHE BOOL "Workarounds")
\ No newline at end of file
diff --git a/user_ops/PointTestCase.py b/user_ops/PointTestCase.py
deleted file mode 100644
index 81c634a..0000000
--- a/user_ops/PointTestCase.py
+++ /dev/null
@@ -1,130 +0,0 @@
-#!/usr/bin/env python
-# -*- coding: utf-8 -*-
-
-# Copyright 2017 ComputerGraphics Tuebingen. All Rights Reserved.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-# ==============================================================================
-# Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
-
-
-import numpy as np
-import tensorflow as tf
-
-np.random.seed(42)
-tf.set_random_seed(42)
-
-
-class FakePointCloud(object):
-    """docstring for FakePointCloud"""
-
-    def __init__(self, B, N, K, Din, Dout, Dp, N2, scaling=1):
-        super(FakePointCloud, self).__init__()
-        assert K < N
-        self.B = B
-        self.N = N
-        self.K = K
-        self.Din = Din
-        self.Dout = Dout
-        self.Dp = Dp
-        self.N2 = N2
-
-    def expected_feature_shape(self):
-        return [self.B, self.Din, self.N]
-
-    def expected_output_shape(self):
-        return [self.B, self.Dout, self.N]
-
-
-def random_values(shape, human_readable=False):
-    """Return random values within range [-10, 10] and precision 2
-    """
-    length = np.prod(shape)
-    return np.arange(length).astype(np.float32).reshape(shape) / float(length)
-
-
-def summary(numeric_grad, graph_grad, name, eps=0.001, max_outputs=20):
-    a, b = numeric_grad.flatten(), graph_grad.flatten()
-    print("summary: %s" % name)
-    print("\ttheirs\t\tours\t\tabs-diff")
-    for i in range(np.prod(numeric_grad.shape)):
-        if np.abs(a[i] - b[i]) > eps and max_outputs > 0:
-            print('%i\t%f\t%f\t%f' % (i, a[i], b[i], np.abs(a[i] - b[i])))
-            max_outputs -= 1
-    if max_outputs == 20:
-        for i in range(max_outputs):
-                print('%i\t%f\t%f\t%f' % (i, a[i], b[i], np.abs(a[i] - b[i])))
-    # print( np.stack([numeric_grad, graph_grad], axis=-1)
-    print("%s - abs-diff (sum): " % name, np.abs(
-        graph_grad - numeric_grad).sum())
-    print("%s - abs-diff (max): " % name, np.abs(
-        graph_grad - numeric_grad).max())
-    print("%s - abs-diff (mean): " % name, np.abs(
-        graph_grad - numeric_grad).mean())
-
-
-# TestPointCloud(B, N, K, Din, Dout, Dp, N2)
-# TPC = FakePointCloud(2, 32, 16, 5, 6, 3, 16)
-# TPC = FakePointCloud(2, 16, 8, 5, 6, 3, 8)
-TPC = FakePointCloud(2, 32, 4, 2, 6, 3, 16)
-# TPC = FakePointCloud(2, 64, 8, 1, 6, 3, 64)
-
-
-class PointTestCase(tf.test.TestCase):
-
-    def __init__(self, methodName="runTest", data=None):
-
-        if data is None:
-            data = TPC
-
-        self.position = random_values([data.B, data.Dp, data.N])
-        self.features = random_values([data.B, data.Din, data.N])
-
-        # make sure, each neighbor hood has no duplicates and first entry is point n
-        # THIS IS IMPORTANT!!
-        self.neighborhood = np.zeros((data.B, data.K, data.N), dtype=np.int32)
-        for b in range(data.B):
-            for n in range(data.N):
-                x = np.arange(data.N)
-                # does not support axis, hence the loop
-                np.random.shuffle(x)
-                offset = np.argwhere(x == n)[0][0]
-                # roll array such that n is first entry
-                x = np.roll(x, -offset)
-                self.neighborhood[b, :, n] = x[:data.K].astype(np.int32)
-
-        self.neighborhood_ds = np.zeros((data.B, data.K, data.N2), dtype=np.int32)
-        for b in range(data.B):
-            for n in range(data.N2):
-                x = np.arange(data.N2)
-                # does not support axis, hence the loop
-                np.random.shuffle(x)
-                offset = np.argwhere(x == n)[0][0]
-                # roll array such that n is first entry
-                x = np.roll(x, -offset)
-                self.neighborhood[b, :, n] = x[:data.K].astype(np.int32)
-
-        self.theta = random_values([1, data.Dp, data.Din, data.Dout])
-        self.bias = random_values([data.Din, data.Dout])
-
-        super(PointTestCase, self).__init__(methodName)
-
-    def init_ops(self):
-        # needs to be called in each method, otherwise graph is empty
-        # probably tf.reset_graph between calls
-        self.features_op = tf.convert_to_tensor(self.features)
-        self.position_op = tf.convert_to_tensor(self.position)
-        self.neighborhood_op = tf.convert_to_tensor(self.neighborhood)
-        self.neighborhood_ds_op = tf.convert_to_tensor(self.neighborhood_ds)
-        self.theta_op = tf.convert_to_tensor(self.theta)
-        self.bias_op = tf.convert_to_tensor(self.bias)
diff --git a/user_ops/__init__.py b/user_ops/__init__.py
index 78449c8..1d80fc4 100644
--- a/user_ops/__init__.py
+++ b/user_ops/__init__.py
@@ -19,27 +19,38 @@ from __future__ import division
 from __future__ import print_function
 
 import tensorflow as tf
+import os
 from tensorflow.python.framework import ops
 from tensorflow.contrib.util import loader
-from tensorflow.python.platform import resource_loader
 
-_flex_convolution_op_so = loader.load_op_library(
-    resource_loader.get_path_to_datafile("flex_conv_op.so"))
-
-_flex_pooling_op_so = loader.load_op_library(
-    resource_loader.get_path_to_datafile("flex_pool_op.so"))
-
-_flex_deconvolution_op_so = loader.load_op_library(
-    resource_loader.get_path_to_datafile("flex_deconv_op.so"))
+__all__ = []
 
 
-# undocumented version
-flex_conv = _flex_convolution_op_so.flex_conv
-flex_conv_grad = _flex_convolution_op_so.flex_conv_grad
-flex_pool = _flex_pooling_op_so.flex_pool
-flex_pool_grad = _flex_pooling_op_so.flex_pool_grad
-flex_deconv = _flex_deconvolution_op_so.flex_deconv
-flex_deconv_grad = _flex_deconvolution_op_so.flex_deconv_grad
+def load_op(name, has_grad=False, public=False):
+  global __all__
+  path = os.path.join(os.path.dirname(__file__), '%s_op.so' % name)
+  if os.path.isfile(path):
+    _module = loader.load_op_library(path)
+    if has_grad:
+      if public:
+        __all__.append('%s' % name)
+        __all__.append('%s_grad' % name)
+      return getattr(_module, '%s' % name), getattr(_module, '%s_grad' % name)
+    else:
+      if public:
+        __all__.append('%s' % name)
+      return getattr(_module, '%s' % name)
+  else:
+    print('[WARNING]: %s does not exists' % name)
+
+knn_bruteforce = load_op('knn_bruteforce', has_grad=False, public=True)
+
+_flex_conv, _flex_conv_grad = load_op(
+    'flex_conv', has_grad=True, public=False)
+_flex_pool, _flex_pool_grad = load_op(
+    'flex_pool', has_grad=True, public=False)
+_flex_deconv, _flex_deconv_grad = load_op(
+    'flex_deconv', has_grad=True, public=False)
 
 
 # pylint: disable=redefined-builtin
@@ -69,7 +80,10 @@ def flex_convolution(features,
   """
 
   with ops.name_scope(name, "flex_convolution"):
-    return flex_conv(features, theta, bias, neighborhood, position)
+    return _flex_conv(features, theta, bias, neighborhood, position)
+
+
+__all__.append('flex_convolution')
 
 
 @ops.RegisterGradient("FlexConv")
@@ -81,7 +95,7 @@ def _FlexConvGrad(op, *grads):  # noqa
   positions = ops.convert_to_tensor(op.inputs[4])
   topdiff = ops.convert_to_tensor(grads[0])
 
-  df, dt, db = flex_conv_grad(
+  df, dt, db = _flex_conv_grad(
       features, theta, bias, neighborhood, positions, topdiff)
 
   df = ops.convert_to_tensor(df, name='gradient_features')
@@ -112,7 +126,10 @@ def flex_pooling(features,
   """
 
   with ops.name_scope(name, "flex_pooling"):
-    return flex_pool(features, neighborhood)
+    return _flex_pool(features, neighborhood)
+
+
+__all__.append('flex_pooling')
 
 
 @ops.RegisterGradient("FlexPool")
@@ -122,7 +139,7 @@ def _FlexPoolGrad(op, *grads):  # noqa
   argmax = ops.convert_to_tensor(op.outputs[1])
   topdiff = ops.convert_to_tensor(grads[0])
 
-  df = flex_pool_grad(features, neighborhood, topdiff, argmax)
+  df = _flex_pool_grad(features, neighborhood, topdiff, argmax)
   df = ops.convert_to_tensor(df, name='gradient_features')
 
   return [df, None]
@@ -153,7 +170,10 @@ def flex_convolution_transpose(features,
   """
 
   with ops.name_scope(name, "flex_convolution_transpose"):
-    return flex_deconv(features, theta, bias, neighborhood, position)
+    return _flex_deconv(features, theta, bias, neighborhood, position)
+
+
+__all__.append('_flex_deconv')
 
 
 @ops.RegisterGradient("FlexDeconv")
@@ -165,7 +185,7 @@ def _FlexDeconvGrad(op, *grads):  # noqa
   positions = ops.convert_to_tensor(op.inputs[4])
   topdiff = ops.convert_to_tensor(grads[0])
 
-  df, dt, db = flex_deconv_grad(
+  df, dt, db = _flex_deconv_grad(
       features, theta, bias, neighborhood, positions, topdiff)
 
   df = ops.convert_to_tensor(df, name='gradient_features')
diff --git a/user_ops/kernels/cuda_utils.h b/user_ops/kernels/cuda_utils.h
new file mode 100644
index 0000000..2b7fc4a
--- /dev/null
+++ b/user_ops/kernels/cuda_utils.h
@@ -0,0 +1,10 @@
+#ifndef LIB_CUDA_UTILS_H_
+#define LIB_CUDA_UTILS_H_
+
+template <typename T>
+__device__ T* DynamicSharedMemory() {
+  extern __shared__ __align__(sizeof(T)) unsigned char s_shm[];
+  return reinterpret_cast<T*>(s_shm);
+}
+
+#endif  // LIB_CUDA_UTILS_H_
diff --git a/user_ops/kernels/flex_conv_kernel.cc b/user_ops/kernels/flex_conv_kernel.cc
index 5ad7c96..7564a6f 100644
--- a/user_ops/kernels/flex_conv_kernel.cc
+++ b/user_ops/kernels/flex_conv_kernel.cc
@@ -70,6 +70,7 @@ struct FlexConvFunctor<CPUDevice, Dtype> {
 };
 
 template struct FlexConvFunctor<CPUDevice, float>;
+template struct FlexConvFunctor<CPUDevice, double>;
 
 template <typename Dtype>
 struct FlexConvGrad<CPUDevice, Dtype> {
@@ -167,7 +168,7 @@ struct FlexConvGrad<CPUDevice, Dtype> {
 
 // template struct FlexConvGrad<CPUDevice, int>;
 template struct FlexConvGrad<CPUDevice, float>;
-// template struct FlexConvGrad<CPUDevice, double>;
+template struct FlexConvGrad<CPUDevice, double>;
 
 }  // namespace functor
 }  // namespace tensorflow
diff --git a/user_ops/kernels/flex_conv_kernel_gpu.cu.cc b/user_ops/kernels/flex_conv_kernel_gpu.cu.cc
index 3925bbf..113be8b 100644
--- a/user_ops/kernels/flex_conv_kernel_gpu.cu.cc
+++ b/user_ops/kernels/flex_conv_kernel_gpu.cu.cc
@@ -12,7 +12,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
-//Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
+// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
 
 #if GOOGLE_CUDA
 
@@ -20,6 +20,7 @@ limitations under the License.
 
 #include <cub/cub.cuh>
 
+#include "cuda_utils.h"
 #include "flex_conv_op.h"
 #include "tensorflow/core/util/cuda_kernel_helper.h"
 
@@ -57,10 +58,10 @@ struct ForwardKernel {
   }
 
   __device__ __forceinline__ void operator()() const {
-    extern __shared__ Dtype s_shm[];
+    Dtype* s_shm = DynamicSharedMemory<Dtype>();
 
-    Dtype* s_theta = (float*)&s_shm[0];
-    Dtype* s_bias = (float*)&s_shm[Dp * C_Din * C_Dout];
+    Dtype* s_theta = (Dtype*)&s_shm[0];
+    Dtype* s_bias = (Dtype*)&s_shm[Dp * C_Din * C_Dout];
 
     // glob ids
     int b = blockIdx.z;
@@ -281,7 +282,8 @@ struct BackwardFeatureKernel {
   }
 
   __device__ __forceinline__ void operator()() const {
-    extern __shared__ float s_shm[];
+    // extern __shared__ float s_shm[];
+    Dtype* s_shm = DynamicSharedMemory<Dtype>();
 
     int i_n = threadIdx.x;
     int i_din = threadIdx.y;
@@ -354,7 +356,10 @@ struct BackwardFeatureKernel {
             W += s_bias[i_din * C_Dout + dout_inner];
             Dtype value = W * s_topdiff[dout_inner * C_N + i_n];
 
-            atomicAdd(
+            // atomicAdd(
+            //     &d_features_out[b * Din * N + din * N + s_nk[k * C_N + i_n]],
+            //     value);
+            tensorflow::CudaAtomicAdd(
                 &d_features_out[b * Din * N + din * N + s_nk[k * C_N + i_n]],
                 value);
           }
@@ -384,6 +389,16 @@ struct BackwardFeatureKernel {
 namespace tensorflow {
 namespace functor {
 
+template <class Dtype, class NBtype>
+struct ForwardKernelType {
+  typedef FlexConvCuda::ForwardKernel<Dtype, NBtype, 3, 128, 32, 32> type;
+};
+
+template <>
+struct ForwardKernelType<float, int> {
+  typedef FlexConvCuda::ForwardKernel<float, int, 3, 128, 32, 64> type;
+};
+
 template <typename Dtype>
 struct FlexConvFunctor<GPUDevice, Dtype> {
   void operator()(::tensorflow::OpKernelContext* ctx, const Tensor& features,
@@ -399,7 +414,11 @@ struct FlexConvFunctor<GPUDevice, Dtype> {
     const int Din = theta.dim_size(2);
     const int Dout = theta.dim_size(3);
 
-    FlexConvCuda::ForwardKernel<Dtype, NBtype, 3, 128, 32, 64> fwk;
+    //    printf("<f> test: %s\n", __PRETTY_FUNCTION__);
+
+    typedef typename ForwardKernelType<Dtype, NBtype>::type FKT;
+
+    FKT fwk;
     fwk.N = N;
     fwk.K = K;
     fwk.Din = Din;
@@ -422,6 +441,7 @@ struct FlexConvFunctor<GPUDevice, Dtype> {
 };
 
 template struct FlexConvFunctor<GPUDevice, float>;
+template struct FlexConvFunctor<GPUDevice, double>;
 
 template <typename Dtype>
 struct FlexConvGrad<GPUDevice, Dtype> {
@@ -459,7 +479,7 @@ struct FlexConvGrad<GPUDevice, Dtype> {
 
     const Dtype* topdiff_ptr = reinterpret_cast<const Dtype*>(topdiff.data());
 
-    Dtype* grad_features_ptr = reinterpret_cast<float*>(grad_features.data());
+    Dtype* grad_features_ptr = reinterpret_cast<Dtype*>(grad_features.data());
     Dtype* grad_theta_ptr = reinterpret_cast<Dtype*>(grad_theta.data());
     Dtype* grad_bias_ptr = reinterpret_cast<Dtype*>(grad_bias.data());
 
@@ -525,6 +545,7 @@ struct FlexConvGrad<GPUDevice, Dtype> {
 };
 
 template struct FlexConvGrad<GPUDevice, float>;
+template struct FlexConvGrad<GPUDevice, double>;
 
 }  // namespace functor
 }  // namespace tensorflow
diff --git a/user_ops/kernels/flex_conv_op.cc b/user_ops/kernels/flex_conv_op.cc
index 2b47b0b..a6b4c4f 100644
--- a/user_ops/kernels/flex_conv_op.cc
+++ b/user_ops/kernels/flex_conv_op.cc
@@ -12,7 +12,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
-//Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
+// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
 
 #include "flex_conv_op.h"
 
@@ -91,32 +91,22 @@ class FlexConvGradOp : public OpKernel {
   }
 };
 
-// Register the CPU kernels.
-#define REGISTER_FLEXCONV_OP_CPU(T)                                   \
-  REGISTER_KERNEL_BUILDER(                                            \
-      Name("FlexConv").Device(DEVICE_CPU).TypeConstraint<T>("T"),     \
-      FlexConvOp<CPUDevice, T>)                                       \
-  REGISTER_KERNEL_BUILDER(                                            \
-      Name("FlexConvGrad").Device(DEVICE_CPU).TypeConstraint<T>("T"), \
-      FlexConvGradOp<CPUDevice, T>)
+#define REGISTER_CUSTOM_OP(NAME, DEVICE, T)                       \
+  REGISTER_KERNEL_BUILDER(                                        \
+      Name(#NAME).Device(DEVICE_##DEVICE).TypeConstraint<T>("T"), \
+      NAME##Op<DEVICE##Device, T>)
 
-TF_CALL_float(REGISTER_FLEXCONV_OP_CPU);
-#undef REGISTER_FLEXCONV_OP_CPU
+REGISTER_CUSTOM_OP(FlexConv, CPU, float);
+REGISTER_CUSTOM_OP(FlexConvGrad, CPU, float);
+REGISTER_CUSTOM_OP(FlexConv, CPU, double);
+REGISTER_CUSTOM_OP(FlexConvGrad, CPU, double);
 
-// Register the GPU kernels.
-// #ifdef GOOGLE_CUDA
-
-#define REGISTER_FLEXCONV_OP_GPU(T)                                   \
-  REGISTER_KERNEL_BUILDER(                                            \
-      Name("FlexConv").Device(DEVICE_GPU).TypeConstraint<T>("T"),     \
-      FlexConvOp<GPUDevice, T>)                                       \
-  REGISTER_KERNEL_BUILDER(                                            \
-      Name("FlexConvGrad").Device(DEVICE_GPU).TypeConstraint<T>("T"), \
-      FlexConvGradOp<GPUDevice, T>)
-
-TF_CALL_float(REGISTER_FLEXCONV_OP_GPU);
-#undef REGISTER_FLEXCONV_OP_GPU
-
-// #endif  // GOOGLE_CUDA
+#ifdef GOOGLE_CUDA
+REGISTER_CUSTOM_OP(FlexConv, GPU, float);
+REGISTER_CUSTOM_OP(FlexConvGrad, GPU, float);
+REGISTER_CUSTOM_OP(FlexConv, GPU, double);
+REGISTER_CUSTOM_OP(FlexConvGrad, GPU, double);
+#endif  // GOOGLE_CUDA
+#undef REGISTER_CUSTOM_OP
 
 }  // namespace tensorflow
diff --git a/user_ops/kernels/flex_deconv_kernel.cc b/user_ops/kernels/flex_deconv_kernel.cc
index d4c618c..32d7358 100644
--- a/user_ops/kernels/flex_deconv_kernel.cc
+++ b/user_ops/kernels/flex_deconv_kernel.cc
@@ -71,6 +71,7 @@ struct FlexDeconvFunctor<CPUDevice, Dtype> {
 };
 
 template struct FlexDeconvFunctor<CPUDevice, float>;
+template struct FlexDeconvFunctor<CPUDevice, double>;
 
 template <typename Dtype>
 struct FlexDeconvGrad<CPUDevice, Dtype> {
@@ -165,9 +166,8 @@ struct FlexDeconvGrad<CPUDevice, Dtype> {
   }
 };
 
-// template struct FlexDeconvGrad<CPUDevice, int>;
 template struct FlexDeconvGrad<CPUDevice, float>;
-// template struct FlexDeconvGrad<CPUDevice, double>;
+template struct FlexDeconvGrad<CPUDevice, double>;
 
 }  // namespace functor
 }  // namespace tensorflow
diff --git a/user_ops/kernels/flex_deconv_kernel_gpu.cu.cc b/user_ops/kernels/flex_deconv_kernel_gpu.cu.cc
index 05518a7..ac929d0 100644
--- a/user_ops/kernels/flex_deconv_kernel_gpu.cu.cc
+++ b/user_ops/kernels/flex_deconv_kernel_gpu.cu.cc
@@ -177,6 +177,7 @@ struct FlexDeconvFunctor<GPUDevice, Dtype> {
 };
 
 template struct FlexDeconvFunctor<GPUDevice, float>;
+template struct FlexDeconvFunctor<GPUDevice, double>;
 
 template <typename Dtype>
 struct FlexDeconvGrad<GPUDevice, Dtype> {
@@ -220,6 +221,7 @@ struct FlexDeconvGrad<GPUDevice, Dtype> {
 };
 
 template struct FlexDeconvGrad<GPUDevice, float>;
+template struct FlexDeconvGrad<GPUDevice, double>;
 
 }  // namespace functor
 }  // namespace tensorflow
diff --git a/user_ops/kernels/flex_deconv_op.cc b/user_ops/kernels/flex_deconv_op.cc
index 05e5d64..424b00f 100644
--- a/user_ops/kernels/flex_deconv_op.cc
+++ b/user_ops/kernels/flex_deconv_op.cc
@@ -12,7 +12,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
-//Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
+// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
 
 #include "flex_deconv_op.h"
 
@@ -88,16 +88,22 @@ class FlexDeconvGradOp : public OpKernel {
   }
 };
 
-#define OPNAME(NAME) NAME##Op
-#define REGISTER(NAME, Dtype)                                    \
-  REGISTER_KERNEL_BUILDER(                                       \
-      Name(#NAME).Device(DEVICE_CPU).TypeConstraint<Dtype>("T"), \
-      OPNAME(NAME) < CPUDevice, Dtype >);                        \
-  REGISTER_KERNEL_BUILDER(                                       \
-      Name(#NAME).Device(DEVICE_GPU).TypeConstraint<Dtype>("T"), \
-      OPNAME(NAME) < GPUDevice, Dtype >);
+#define REGISTER_CUSTOM_OP(NAME, DEVICE, T)                       \
+  REGISTER_KERNEL_BUILDER(                                        \
+      Name(#NAME).Device(DEVICE_##DEVICE).TypeConstraint<T>("T"), \
+      NAME##Op<DEVICE##Device, T>)
 
-REGISTER(FlexDeconv, float);
-REGISTER(FlexDeconvGrad, float);
+REGISTER_CUSTOM_OP(FlexDeconv, CPU, float);
+REGISTER_CUSTOM_OP(FlexDeconvGrad, CPU, float);
+REGISTER_CUSTOM_OP(FlexDeconv, CPU, double);
+REGISTER_CUSTOM_OP(FlexDeconvGrad, CPU, double);
+
+#ifdef GOOGLE_CUDA
+REGISTER_CUSTOM_OP(FlexDeconv, GPU, float);
+REGISTER_CUSTOM_OP(FlexDeconvGrad, GPU, float);
+REGISTER_CUSTOM_OP(FlexDeconv, GPU, double);
+REGISTER_CUSTOM_OP(FlexDeconvGrad, GPU, double);
+#endif  // GOOGLE_CUDA
+#undef REGISTER_CUSTOM_OP
 
 }  // namespace tensorflow
diff --git a/user_ops/kernels/flex_pool_kernel.cc b/user_ops/kernels/flex_pool_kernel.cc
index 4158286..31ef5c1 100644
--- a/user_ops/kernels/flex_pool_kernel.cc
+++ b/user_ops/kernels/flex_pool_kernel.cc
@@ -12,7 +12,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
-//Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
+// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
 
 #include "flex_pool_op.h"
 #include "tensorflow/core/framework/op.h"
@@ -59,6 +59,7 @@ struct FlexPoolFunctor<CPUDevice, Dtype> {
 };
 
 template struct FlexPoolFunctor<CPUDevice, float>;
+template struct FlexPoolFunctor<CPUDevice, double>;
 
 template <typename Dtype>
 struct FlexPoolGrad<CPUDevice, Dtype> {
@@ -94,9 +95,8 @@ struct FlexPoolGrad<CPUDevice, Dtype> {
   }
 };
 
-// template struct FlexPoolGrad<CPUDevice, int>;
 template struct FlexPoolGrad<CPUDevice, float>;
-// template struct FlexPoolGrad<CPUDevice, double>;
+template struct FlexPoolGrad<CPUDevice, double>;
 
 }  // namespace functor
 }  // namespace tensorflow
diff --git a/user_ops/kernels/flex_pool_kernel_gpu.cu.cc b/user_ops/kernels/flex_pool_kernel_gpu.cu.cc
index ecb39d2..8ad9bbe 100644
--- a/user_ops/kernels/flex_pool_kernel_gpu.cu.cc
+++ b/user_ops/kernels/flex_pool_kernel_gpu.cu.cc
@@ -12,7 +12,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
-//Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
+// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
 
 #if GOOGLE_CUDA
 
@@ -30,7 +30,7 @@ inline int up2(int len, int th) { return (len - 1) / th + 1; }
 template <typename Dtype>
 __global__ void forward(const int B, const int N, const int K, const int D,
                         const Dtype* features, const int* neighborhood,
-                        Dtype* output, int* argmax, float float_min_value) {
+                        Dtype* output, int* argmax, Dtype float_min_value) {
   // features: each feature description for each point [B, D, N].
   // neighborhood: all K nearest neighbors [B, K, N].
   // output: each feature description for each point [B, D, N].
@@ -41,14 +41,14 @@ __global__ void forward(const int B, const int N, const int K, const int D,
        d += blockDim.y * gridDim.y) {
     for (int n = blockIdx.x * blockDim.x + threadIdx.x; n < N;
          n += blockDim.x * gridDim.x) {
-      float best_value = float_min_value;
+      Dtype best_value = float_min_value;
       int best_id = 0;
 
       const int current_flat = b * D * N + d * N + n;
 
       for (int k_ = 0; k_ < K; ++k_) {
         const int other_global_id = neighborhood[b * K * N + k_ * N + n];
-        const float v = features[b * D * N + d * N + other_global_id];
+        const Dtype v = features[b * D * N + d * N + other_global_id];
 
         if (best_value < v) {
           best_id = other_global_id;
@@ -113,13 +113,9 @@ struct FlexPoolFunctor<GPUDevice, Dtype> {
     dim3 grid(up2(N, threads), up2(D, threads), B);
 
     forward<Dtype><<<grid, block>>>(
-        B, N, K, D,
-
-        features_.flat<Dtype>().data(), neighborhood_.flat<int>().data(),
-
-        output_->flat<Dtype>().data(), argmax_->flat<int>().data(),
-
-        std::numeric_limits<Dtype>::lowest());
+        B, N, K, D, features_.flat<Dtype>().data(),
+        neighborhood_.flat<int>().data(), output_->flat<Dtype>().data(),
+        argmax_->flat<int>().data(), std::numeric_limits<Dtype>::lowest());
 
     if (!ctx->eigen_gpu_device().ok()) {
       ctx->SetStatus(
@@ -129,6 +125,7 @@ struct FlexPoolFunctor<GPUDevice, Dtype> {
 };
 
 template struct FlexPoolFunctor<GPUDevice, float>;
+template struct FlexPoolFunctor<GPUDevice, double>;
 
 template <typename Dtype>
 struct FlexPoolGrad<GPUDevice, Dtype> {
@@ -149,13 +146,9 @@ struct FlexPoolGrad<GPUDevice, Dtype> {
                grad_features_->NumElements() * sizeof(Dtype));
 
     backward<Dtype><<<grid, block>>>(
-        B, N, K, D,
-
-        features_.flat<Dtype>().data(), neighborhood_.flat<int>().data(),
-
-        topdiff_.flat<Dtype>().data(), argmax_.flat<int>().data(),
-
-        grad_features_->flat<Dtype>().data());
+        B, N, K, D, features_.flat<Dtype>().data(),
+        neighborhood_.flat<int>().data(), topdiff_.flat<Dtype>().data(),
+        argmax_.flat<int>().data(), grad_features_->flat<Dtype>().data());
 
     if (!ctx->eigen_gpu_device().ok()) {
       ctx->SetStatus(tensorflow::errors::Internal("CUDA: FlexPoolGrad Error!"));
@@ -164,6 +157,7 @@ struct FlexPoolGrad<GPUDevice, Dtype> {
 };
 
 template struct FlexPoolGrad<GPUDevice, float>;
+template struct FlexPoolGrad<GPUDevice, double>;
 
 }  // namespace functor
 }  // namespace tensorflow
diff --git a/user_ops/kernels/flex_pool_op.cc b/user_ops/kernels/flex_pool_op.cc
index 508a8dd..2d14185 100644
--- a/user_ops/kernels/flex_pool_op.cc
+++ b/user_ops/kernels/flex_pool_op.cc
@@ -12,7 +12,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 ==============================================================================*/
-//Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
+// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
 
 #include "flex_pool_op.h"
 
@@ -33,7 +33,6 @@ class FlexPoolOp : public OpKernel {
   explicit FlexPoolOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
 
   void Compute(OpKernelContext* ctx) override {
-    // printf("--> Compute CPU Version <--\n");
     const Tensor& features_ = ctx->input(0);
     const Tensor& neighborhood_ = ctx->input(1);
 
@@ -65,7 +64,6 @@ class FlexPoolGradOp : public OpKernel {
   explicit FlexPoolGradOp(OpKernelConstruction* ctx) : OpKernel(ctx) {}
 
   void Compute(OpKernelContext* ctx) override {
-    // printf("--> Compute CPU Version <--\n");
     const Tensor& features_ = ctx->input(0);
     const Tensor& neighborhood_ = ctx->input(1);
     const Tensor& topdiff_ = ctx->input(2);
@@ -82,32 +80,22 @@ class FlexPoolGradOp : public OpKernel {
   }
 };
 
-// Register the CPU kernels.
-#define REGISTER_FLEXPOOL_OP_CPU(T)                                   \
-  REGISTER_KERNEL_BUILDER(                                            \
-      Name("FlexPool").Device(DEVICE_CPU).TypeConstraint<T>("T"),     \
-      FlexPoolOp<CPUDevice, T>)                                       \
-  REGISTER_KERNEL_BUILDER(                                            \
-      Name("FlexPoolGrad").Device(DEVICE_CPU).TypeConstraint<T>("T"), \
-      FlexPoolGradOp<CPUDevice, T>)
+#define REGISTER_CUSTOM_OP(NAME, DEVICE, T)                       \
+  REGISTER_KERNEL_BUILDER(                                        \
+      Name(#NAME).Device(DEVICE_##DEVICE).TypeConstraint<T>("T"), \
+      NAME##Op<DEVICE##Device, T>)
 
-TF_CALL_float(REGISTER_FLEXPOOL_OP_CPU);
-#undef REGISTER_FLEXPOOL_OP_CPU
+REGISTER_CUSTOM_OP(FlexPool, CPU, float);
+REGISTER_CUSTOM_OP(FlexPoolGrad, CPU, float);
+REGISTER_CUSTOM_OP(FlexPool, CPU, double);
+REGISTER_CUSTOM_OP(FlexPoolGrad, CPU, double);
 
-// Register the GPU kernels.
 #ifdef GOOGLE_CUDA
-
-#define REGISTER_FLEXPOOL_OP_GPU(T)                                   \
-  REGISTER_KERNEL_BUILDER(                                            \
-      Name("FlexPool").Device(DEVICE_GPU).TypeConstraint<T>("T"),     \
-      FlexPoolOp<GPUDevice, T>)                                       \
-  REGISTER_KERNEL_BUILDER(                                            \
-      Name("FlexPoolGrad").Device(DEVICE_GPU).TypeConstraint<T>("T"), \
-      FlexPoolGradOp<GPUDevice, T>)
-
-TF_CALL_float(REGISTER_FLEXPOOL_OP_GPU);
-#undef REGISTER_FLEXPOOL_OP_GPU
-
+REGISTER_CUSTOM_OP(FlexPool, GPU, float);
+REGISTER_CUSTOM_OP(FlexPoolGrad, GPU, float);
+REGISTER_CUSTOM_OP(FlexPool, GPU, double);
+REGISTER_CUSTOM_OP(FlexPoolGrad, GPU, double);
 #endif  // GOOGLE_CUDA
+#undef REGISTER_CUSTOM_OP
 
 }  // namespace tensorflow
diff --git a/user_ops/kernels/knn_bruteforce_kernel.cc b/user_ops/kernels/knn_bruteforce_kernel.cc
new file mode 100644
index 0000000..2d846ce
--- /dev/null
+++ b/user_ops/kernels/knn_bruteforce_kernel.cc
@@ -0,0 +1,79 @@
+/* Copyright 2018 ComputerGraphics Tuebingen. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
+
+#include "knn_bruteforce_op.h"
+#include "tensorflow/core/framework/op.h"
+
+namespace tensorflow {
+
+namespace functor {
+
+template <typename Dtype, typename NBtype>
+struct KnnBruteforceFunctor<CPUDevice, Dtype, NBtype> {
+  void operator()(::tensorflow::OpKernelContext* ctx, const Tensor& positions_,
+                  Tensor* neighborhood_out_, Tensor* distances_,
+                  Tensor* timings) {
+    // positions [B, Dp, N]
+    // neighborhood_out_ [B, N, K]
+    // distances_ [B, N, K]
+    const auto positions = positions_.tensor<Dtype, 3>();
+    auto neighborhood_out = neighborhood_out_->tensor<NBtype, 3>();
+    auto distances_out = distances_->tensor<Dtype, 3>();
+
+    const int B = positions_.dim_size(0);
+    const int Dp = positions_.dim_size(1);
+    const int N = positions_.dim_size(2);
+
+    const int K = neighborhood_out_->dim_size(2);
+
+    for (int b = 0; b < B; ++b) {
+      const Dtype* pc_raw = positions_.flat<Dtype>().data() + b * Dp * N;
+      Eigen::Map<const Eigen::Matrix<Dtype, Eigen::Dynamic, Eigen::Dynamic,
+                                     Eigen::RowMajor>>
+          pc(pc_raw, Dp, N);
+      // pc: [Dp, N]
+
+      for (int n = 0; n < N; ++n) {
+        const auto query = pc.col(n);
+        const auto diffs =
+            (pc.colwise() - query).array() * (pc.colwise() - query).array();
+        const auto distances = (diffs.colwise().sum()).array().sqrt();
+
+        std::vector<Dtype> vec_dist;
+        std::vector<Dtype> vec_ids;
+
+        for (int i = 0; i < N; ++i) {
+          vec_dist.push_back(distances(i));
+          vec_ids.push_back(i);
+        }
+
+        std::sort(std::begin(vec_ids), std::end(vec_ids),
+                  [&](int i1, int i2) { return vec_dist[i1] < vec_dist[i2]; });
+
+        for (int k = 0; k < K; ++k) {
+          neighborhood_out(b, n, k) = vec_ids[k];
+          distances_out(b, n, k) = vec_dist[vec_ids[k]];
+        }
+      }
+    }
+  }
+};
+
+template struct KnnBruteforceFunctor<CPUDevice, float, int>;
+template struct KnnBruteforceFunctor<CPUDevice, double, int>;
+
+}  // namespace functor
+}  // namespace tensorflow
diff --git a/user_ops/kernels/knn_bruteforce_kernel_gpu.cu.cc b/user_ops/kernels/knn_bruteforce_kernel_gpu.cu.cc
new file mode 100644
index 0000000..40bbfd5
--- /dev/null
+++ b/user_ops/kernels/knn_bruteforce_kernel_gpu.cu.cc
@@ -0,0 +1,273 @@
+/* Copyright 2018 ComputerGraphics Tuebingen. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
+
+#if GOOGLE_CUDA
+
+#define EIGEN_USE_GPU
+
+#include <vector>
+#include "cuda_utils.h"
+
+#include <cuda.h>
+#include <helper_cuda.h>
+
+#include <cub/cub.cuh>
+
+#include <curand.h>
+#include <curand_kernel.h>
+
+#include <limits>
+
+#include "knn_bruteforce_op.h"
+#include "tensorflow/core/util/cuda_kernel_helper.h"
+
+template <typename Dtype, typename NBtype, int C_THREADS = 256, int C_VPT = 2>
+struct BlockBFKernel;
+
+template <typename Dtype, typename NBtype, int C_THREADS, int C_VPT>
+__global__ void runBlockBFKernel(
+    const BlockBFKernel<Dtype, NBtype, C_THREADS, C_VPT> kernel) {
+  kernel();
+}
+
+template <typename Dtype, typename NBtype, int C_THREADS, int C_VPT>
+struct BlockBFKernel {
+  void launch(int B) {
+    dim3 block(C_THREADS, 1);
+    dim3 grid(N, 1, B);
+
+    size_t shm_size = Dp * sizeof(Dtype);
+
+    runBlockBFKernel<<<grid, block, shm_size>>>((*this));
+  }
+
+  __device__ __forceinline__ void operator()() const {
+    // extern __shared__ Dtype s_shm[];
+    Dtype* s_shm = DynamicSharedMemory<Dtype>();
+    Dtype* s_point = (Dtype*)&s_shm[0];
+
+    if (N > C_THREADS * C_VPT) {
+      printf(
+          "<BlockBFKernel> Critical problem!!!!! Not enough resources spend "
+          "for N points!!! %d < %d \n",
+          C_THREADS * C_VPT, N);
+      return;
+    }
+
+    int b = blockIdx.z;
+    int y = blockIdx.x;
+
+    int tid = threadIdx.x;
+    for (int dpi = tid; dpi < Dp; dpi += blockDim.x) {
+      s_point[dpi] = d_data[b * N * Dp + dpi * N + y];  // not aligned with
+                                                        // data!
+    }
+
+    __syncthreads();
+
+    Dtype thread_dists[C_VPT];  // keys
+    NBtype thread_ids[C_VPT];   // values
+
+    typedef cub::BlockRadixSort<Dtype, C_THREADS, C_VPT, NBtype> BlockRadixSort;
+    typedef cub::BlockStore<Dtype, C_THREADS, C_VPT,
+                            cub::BLOCK_STORE_WARP_TRANSPOSE>
+        BlockStoreDists;
+    typedef cub::BlockStore<NBtype, C_THREADS, C_VPT,
+                            cub::BLOCK_STORE_WARP_TRANSPOSE>
+        BlockStoreIds;
+
+    // Allocate shared memory
+    __shared__ union {
+      typename BlockRadixSort::TempStorage sort;
+      typename BlockStoreDists::TempStorage store_dists;
+      typename BlockStoreIds::TempStorage store_ids;
+
+    } temp_storage;
+
+    for (int vpt_i = 0; vpt_i < C_VPT; ++vpt_i) {
+      int x = vpt_i * C_THREADS + threadIdx.x;
+
+      if (x < N) {
+        Dtype sum = 0.f;
+        for (int dpi = 0; dpi < Dp; ++dpi) {
+          //					Dtype val = d_data[b * N * Dp + dpi * N +
+          //x]
+          //- d_data[b * N * Dp + dpi * N + y];
+          Dtype val = d_data[b * N * Dp + dpi * N + x] - s_point[dpi];
+          sum += val * val;
+        }
+        thread_dists[vpt_i] = sqrt(sum);
+        thread_ids[vpt_i] = x;
+      } else {
+        thread_dists[vpt_i] = std::numeric_limits<Dtype>::max();
+        thread_ids[vpt_i] = -1;
+      }
+    }
+
+    BlockRadixSort(temp_storage.sort).Sort(thread_dists, thread_ids);
+    __syncthreads();
+
+    BlockStoreIds(temp_storage.store_ids)
+        .Store(&d_knn_ids[b * N * K + y * K], thread_ids, K);
+    __syncthreads();
+
+    BlockStoreDists(temp_storage.store_dists)
+        .Store(&d_knn_dists[b * N * K + y * K], thread_dists, K);
+  }
+
+  const Dtype* d_data;
+
+  Dtype* d_knn_dists;
+  NBtype* d_knn_ids;
+
+  int N;
+  int Dp;
+  int K;
+};
+
+template <typename Dtype, typename NBtype>
+struct BlockBFKernelAttributesSetter {
+  template <typename T>
+  void setAttributes(T& kernel) {
+    kernel.d_data = d_data;
+    kernel.d_knn_dists = d_knn_dists;
+    kernel.d_knn_ids = d_knn_ids;
+
+    kernel.N = N;
+    kernel.Dp = Dp;
+    kernel.K = K;
+  }
+
+  const Dtype* d_data;
+
+  Dtype* d_knn_dists;
+  NBtype* d_knn_ids;
+
+  int N;
+  int Dp;
+  int K;
+};
+
+namespace tensorflow {
+namespace functor {
+
+template <typename Dtype, typename NBtype>
+struct KnnBruteforceFunctor<GPUDevice, Dtype, NBtype> {
+  void operator()(::tensorflow::OpKernelContext* ctx, const Tensor& positions,
+                  Tensor* neighborhood_out, Tensor* distances,
+                  Tensor* timings) {
+    //		printf("GPU: KNNBF! \n");
+
+    //		printf("return_timings: %d \n",return_timings);
+
+    const int B = positions.dim_size(0);
+    const int D = positions.dim_size(1);
+    const int N = positions.dim_size(2);
+
+    const int K = neighborhood_out->dim_size(2);
+    //		printf("B: %d | N: %d | K: %d \n", B, N, K);
+
+    BlockBFKernelAttributesSetter<Dtype, NBtype> attr;
+    attr.d_data = positions.flat<Dtype>().data();
+    attr.d_knn_ids = neighborhood_out->flat<NBtype>().data();
+    attr.d_knn_dists = distances->flat<Dtype>().data();
+
+    attr.N = N;
+    attr.Dp = D;
+    attr.K = K;
+
+    cudaEvent_t start, stop;
+    if (return_timings) {
+      cudaEventCreate(&start);
+      cudaEventCreate(&stop);
+      cudaEventRecord(start, 0);
+    }
+
+    // TODO: beautify
+    if (N <= 32) {
+      BlockBFKernel<Dtype, NBtype, 32, 1> knn;
+      attr.setAttributes(knn);
+      knn.launch(B);
+    } else if (N <= 64) {
+      BlockBFKernel<Dtype, NBtype, 64, 1> knn;
+      attr.setAttributes(knn);
+      knn.launch(B);
+    } else if (N <= 128) {
+      BlockBFKernel<Dtype, NBtype, 128, 1> knn;
+      attr.setAttributes(knn);
+      knn.launch(B);
+    } else if (N <= 256) {
+      BlockBFKernel<Dtype, NBtype, 128, 2> knn;
+      attr.setAttributes(knn);
+      knn.launch(B);
+    } else if (N <= 512) {
+      BlockBFKernel<Dtype, NBtype, 128, 4> knn;
+      attr.setAttributes(knn);
+      knn.launch(B);
+    } else if (N <= 1024) {
+      BlockBFKernel<Dtype, NBtype, 256, 4> knn;
+      attr.setAttributes(knn);
+      knn.launch(B);
+    } else if (N <= 2048) {
+      BlockBFKernel<Dtype, NBtype, 256, 8> knn;
+      attr.setAttributes(knn);
+      knn.launch(B);
+    } else if (N <= 4096) {
+      BlockBFKernel<Dtype, NBtype, 512, 8> knn;
+      attr.setAttributes(knn);
+      knn.launch(B);
+    } else if (N <= 1024 * 8) {
+      BlockBFKernel<Dtype, NBtype, 1024, 8> knn;
+      attr.setAttributes(knn);
+      knn.launch(B);
+    } else {
+      printf(
+          "point sets greater then 8k are note yet supported!! Change to "
+          "knn_graph operation!! \n");
+    }
+
+    if (return_timings) {
+      float time;
+      cudaEventRecord(stop, 0);
+      cudaEventSynchronize(stop);
+      cudaEventElapsedTime(&time, start, stop);
+      //			printf("complete elapsed time for KNNBf: %f ms
+      //\n", time);
+      Dtype time2 = time;
+
+      cudaMemcpy(timings->flat<Dtype>().data(), &time2, sizeof(Dtype),
+                 cudaMemcpyHostToDevice);
+
+      cudaEventDestroy(start);
+      cudaEventDestroy(stop);
+    }
+
+    cudaDeviceSynchronize();
+    getLastCudaError("KNNBF execution failed");
+    checkCudaErrors(cudaDeviceSynchronize());
+  }
+
+  bool return_timings;
+};
+
+template struct KnnBruteforceFunctor<GPUDevice, float, int>;
+// // too much shared memory
+// template struct KnnBruteforceFunctor<GPUDevice, double, int>;
+
+}  // namespace functor
+}  // namespace tensorflow
+
+#endif  // GOOGLE_CUDA
diff --git a/user_ops/kernels/knn_bruteforce_op.cc b/user_ops/kernels/knn_bruteforce_op.cc
new file mode 100644
index 0000000..7e9415b
--- /dev/null
+++ b/user_ops/kernels/knn_bruteforce_op.cc
@@ -0,0 +1,81 @@
+/* Copyright 2018 ComputerGraphics Tuebingen. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
+
+#include "knn_bruteforce_op.h"
+
+#include <stdio.h>
+#include <type_traits>
+
+#include "tensorflow/core/framework/op.h"
+#include "tensorflow/core/framework/op_kernel.h"
+
+namespace tensorflow {
+
+// Forward-Pass (CPU, GPU)
+// --------------------------------------------------
+template <typename Device, typename Dtype>
+class KnnBruteforceOp : public OpKernel {
+ public:
+  explicit KnnBruteforceOp(OpKernelConstruction* ctx) : OpKernel(ctx) {
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("K", &K));
+    OP_REQUIRES_OK(ctx, ctx->GetAttr("return_timings", &return_timings));
+  }
+
+  void Compute(OpKernelContext* ctx) override {
+    const Tensor& positions = ctx->input(0);
+    const Tensor& neighborhood_in = ctx->input(1);
+
+    const int B = positions.shape().dim_size(0);
+    const int N = positions.shape().dim_size(2);
+
+    Tensor* neighborhood_out = nullptr;
+    Tensor* distances = nullptr;
+    Tensor* timings = nullptr;
+
+    OP_REQUIRES_OK(ctx, ctx->allocate_output(0, TensorShape({B, N, K}),
+                                             &neighborhood_out));
+
+    OP_REQUIRES_OK(ctx,
+                   ctx->allocate_output(1, TensorShape({B, N, K}), &distances));
+
+    OP_REQUIRES_OK(ctx, ctx->allocate_output(2, TensorShape({1}), &timings));
+
+    ::tensorflow::functor::KnnBruteforceFunctor<Device, Dtype, int> knnBFF;
+    knnBFF.return_timings = return_timings;
+    knnBFF(ctx, positions, neighborhood_out, distances, timings);
+  }
+
+ private:
+  TF_DISALLOW_COPY_AND_ASSIGN(KnnBruteforceOp);
+
+  int K;
+  bool return_timings;
+  //	int subBatch;
+};
+
+#define REGISTER_CUSTOM_OP(NAME, DEVICE, T)                       \
+  REGISTER_KERNEL_BUILDER(                                        \
+      Name(#NAME).Device(DEVICE_##DEVICE).TypeConstraint<T>("T"), \
+      NAME##Op<DEVICE##Device, T>)
+
+REGISTER_CUSTOM_OP(KnnBruteforce, CPU, float);
+
+#ifdef GOOGLE_CUDA
+REGISTER_CUSTOM_OP(KnnBruteforce, GPU, float);
+#endif  // GOOGLE_CUDA
+#undef REGISTER_CUSTOM_OP
+
+}  // namespace tensorflow
diff --git a/user_ops/kernels/knn_bruteforce_op.h b/user_ops/kernels/knn_bruteforce_op.h
new file mode 100644
index 0000000..7755c87
--- /dev/null
+++ b/user_ops/kernels/knn_bruteforce_op.h
@@ -0,0 +1,44 @@
+/* Copyright 2018 ComputerGraphics Tuebingen. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+// Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
+
+#ifndef LIB_KNN_BF_OP_H_
+#define LIB_KNN_BF_OP_H_
+
+#include "tensorflow/core/framework/op_kernel.h"
+
+namespace tensorflow {
+class OpKernelContext;
+class Tensor;
+
+using CPUDevice = Eigen::ThreadPoolDevice;
+using GPUDevice = Eigen::GpuDevice;
+}  // namespace tensorflow
+
+namespace tensorflow {
+namespace functor {
+
+template <typename Device, typename Dtype, typename NBtype>
+struct KnnBruteforceFunctor {
+  void operator()(::tensorflow::OpKernelContext *ctx, const Tensor &positions,
+                  Tensor *neighborhood_out, Tensor *distances, Tensor *timings);
+
+  bool return_timings;
+};
+
+}  // namespace functor
+}  // namespace tensorflow
+
+#endif  // LIB_KNN_BF_OP_H_
diff --git a/user_ops/misc.py b/user_ops/misc.py
new file mode 100644
index 0000000..afddddd
--- /dev/null
+++ b/user_ops/misc.py
@@ -0,0 +1,128 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+# Copyright 2017 ComputerGraphics Tuebingen. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
+
+
+import numpy as np
+import tensorflow as tf
+from tabulate import tabulate
+from scipy.spatial.distance import pdist, squareform
+
+np.random.seed(42)
+tf.set_random_seed(42)
+
+
+class FakePointCloud(object):
+  """docstring for FakePointCloud"""
+
+  def __init__(self, B, N, K, Din, Dout, Dp, N2=1, scaling=1):
+    super(FakePointCloud, self).__init__()
+    assert K < N
+    self.B = B
+    self.N = N
+    self.K = K
+    self.Din = Din
+    self.Dout = Dout
+    self.Dp = Dp
+    self.N2 = N2
+
+    dtype = np.float64
+
+    def find_neighbors(positions, K):
+      # B, Dpos, N
+      all_neighbors = []
+      for batch in positions:
+        distances = squareform(pdist(batch.T, 'euclidean'))
+        all_neighbors.append(np.argsort(distances, axis=1)[:, :K])
+      return np.array(all_neighbors).transpose(0, 2, 1)
+
+    def random_values(shape):
+      # return (np.random.randn(*shape) * 100).astype(np.int32).astype(np.float32)
+      # return (np.random.randn(*shape) * 100).astype(np.int32).astype(np.float32)
+      return np.random.randn(*shape).astype(np.float32)
+
+    self.theta = random_values([1, self.Dp, self.Din, self.Dout]).astype(dtype)
+    self.bias = random_values([self.Din, self.Dout]).astype(dtype)
+
+    self.position = random_values([self.B, self.Dp, self.N]).astype(dtype)
+    self.features = random_values([self.B, self.Din, self.N]).astype(dtype)
+    self.neighborhood = find_neighbors(
+        self.position, self.K).astype(dtype=np.int32)
+
+  def init_ops(self, dtype=np.float32):
+    self.theta_op = tf.convert_to_tensor(self.theta.astype(dtype))
+    self.bias_op = tf.convert_to_tensor(self.bias.astype(dtype))
+
+    self.features_op = tf.convert_to_tensor(self.features.astype(dtype))
+    self.position_op = tf.convert_to_tensor(self.position.astype(dtype))
+    self.neighborhood_op = tf.convert_to_tensor(self.neighborhood)
+
+  def expected_feature_shape(self):
+    return [self.B, self.Din, self.N]
+
+  def expected_output_shape(self):
+    return [self.B, self.Dout, self.N]
+
+
+class VerboseTestCase(tf.test.TestCase):
+
+  def assertAllClose(self, a, b, rtol=1e-6, atol=1e-6):
+    max_outputs = 20
+
+    def max_tol(b):
+      return atol + rtol * np.abs(b)
+
+    if not np.allclose(a, b, rtol=rtol, atol=atol):
+      cond = np.logical_or(
+          np.abs(a - b) > atol + rtol * np.abs(b),
+          np.isnan(a) != np.isnan(b))
+
+      lines = []
+      if a.ndim:
+
+        shape = a.shape
+
+        a = a.flatten()
+        b = b.flatten()
+
+        cond = np.logical_or(
+            np.abs(a - b) > atol + rtol * np.abs(b),
+            np.isnan(a) != np.isnan(b))
+
+        idxArr = np.arange(a.shape[0])[np.where(cond)]
+        xArr = a[np.where(cond)]
+        yArr = b[np.where(cond)]
+        for idx, x, y in zip(idxArr, xArr, yArr):
+          idx = np.unravel_index(idx, shape)
+          lines.append((idx, x, y, np.abs(x - y), max_tol(y)))
+          max_outputs -= 1
+          if max_outputs == 0:
+            break
+        print(tabulate(lines, headers=["index", "actual", "expected",
+                                       "diff", "max-tol"]))
+        print("diff (sum): ", np.abs(a - b).sum())
+        print("diff (max): ", np.abs(a - b).max())
+        print("diff (mean): ", np.abs(a - b).mean())
+      else:
+        # np.where is broken for scalars
+        x, y = a, b
+        lines.append((x, y, np.abs(x - y), max_tol(y)))
+        print(tabulate(lines, headers=["actual", "expected",
+                                       "diff", "max-tol"]))
+
+      assert np.allclose(a, b, rtol=rtol, atol=atol, equal_nan=True), "failed"
diff --git a/user_ops/ops/knn_bruteforce.cc b/user_ops/ops/knn_bruteforce.cc
new file mode 100644
index 0000000..23c688f
--- /dev/null
+++ b/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.
diff --git a/user_ops/test_all.py b/user_ops/test_all.py
new file mode 100644
index 0000000..2284a8f
--- /dev/null
+++ b/user_ops/test_all.py
@@ -0,0 +1,29 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+# Copyright 2017 ComputerGraphics Tuebingen. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
+
+import unittest
+import glob
+
+operations = [fn[:-3] for fn in glob.glob('test_*.py') if not '_all' in fn]
+suite = unittest.TestSuite()
+
+for op in operations:
+  suite.addTest(unittest.defaultTestLoader.loadTestsFromName(op))
+
+unittest.TextTestRunner().run(suite)
diff --git a/user_ops/test_flex_convolution.py b/user_ops/test_flex_convolution.py
index 4465474..12d2262 100644
--- a/user_ops/test_flex_convolution.py
+++ b/user_ops/test_flex_convolution.py
@@ -18,91 +18,119 @@
 # Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
 
 
-from PointTestCase import TPC, PointTestCase, summary
+from misc import FakePointCloud, VerboseTestCase
 import tensorflow as tf
-
+import numpy as np
 from __init__ import flex_convolution
 
+case = FakePointCloud(B=2, N=32, K=4, Din=2, Dout=6, Dp=3)
 
-class FlexConvTest(PointTestCase):
+
+class FlexConvTest(VerboseTestCase):
   def __init__(self, methodName="runTest"):
     super(FlexConvTest, self).__init__(methodName)
 
-  def _forward(self, use_gpu=False, force_gpu=False):
-    self.init_ops()
+  def _forward(self, use_gpu=False, force_gpu=False, dtype=np.float32):
+    case.init_ops(dtype=dtype)
     with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu) as sess:
-      actual_op = flex_convolution(self.features_op,
-                                   self.position_op, self.neighborhood_op,
-                                   self.theta_op, self.bias_op)
+      actual_op = flex_convolution(case.features_op,
+                                   case.position_op, case.neighborhood_op,
+                                   case.theta_op, case.bias_op)
       actual = sess.run(actual_op)
     return actual
 
-  def test_forward(self):
+  def test_forward(self, dtype=np.float32):
     cpu = self._forward(use_gpu=False)
     gpu = self._forward(use_gpu=True)
-    self.assertAllClose(cpu, gpu, 1e-5, 1e-5)
+    self.assertAllClose(cpu, gpu, 1e-4)
 
-  def _backward_features(self, use_gpu=False):
-    self.init_ops()
-    with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu):
-      actual_op = flex_convolution(self.features_op,
-                                   self.position_op, self.neighborhood_op,
-                                   self.theta_op, self.bias_op)
-      graph_features_grad, num_features_grad = tf.test.compute_gradient(
-          [self.features_op], [self.features.shape], actual_op,
-          TPC.expected_output_shape())[0]
-      summary(num_features_grad, graph_features_grad, 'self.features')
+  def test_forward_features_gpu_floats(self):
+    cpu32 = self._forward(use_gpu=True, dtype=np.float32)
+    cpu64 = self._forward(use_gpu=True, dtype=np.float64)
+    self.assertAllClose(cpu32, cpu64)
 
-      err = tf.test.compute_gradient_error([self.features_op],
-                                           [self.features.shape],
-                                           actual_op, TPC.expected_output_shape())
-      self.assertLess(err, 1e-2)
+  def _backward_features(self, use_gpu=False, dtype=np.float32, numdiff=True):
+    case.init_ops(dtype=dtype)
+    with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu) as sess:
+      actual_op = flex_convolution(case.features_op,
+                                   case.position_op, case.neighborhood_op,
+                                   case.theta_op, case.bias_op)
+      if numdiff:
+        return tf.test.compute_gradient(
+            [case.features_op], [case.features.shape], actual_op,
+            case.expected_output_shape())[0]
+      else:
+        return sess.run(tf.gradients(actual_op, [case.features_op]))[0]
 
-  def _backward_bias(self, use_gpu=False):
-    self.init_ops()
-    with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu):
-      actual_op = flex_convolution(self.features_op,
-                                   self.position_op, self.neighborhood_op,
-                                   self.theta_op, self.bias_op)
+  def _backward_bias(self, use_gpu=False, dtype=np.float32, numdiff=True):
+    case.init_ops(dtype=dtype)
+    with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu) as sess:
+      actual_op = flex_convolution(case.features_op,
+                                   case.position_op, case.neighborhood_op,
+                                   case.theta_op, case.bias_op)
 
-      graph_bias_grad, num_bias_grad = tf.test.compute_gradient(
-          [self.bias_op], [self.bias.shape], actual_op,
-          TPC.expected_output_shape())[0]
-      summary(num_bias_grad, graph_bias_grad, 'self.bias')
+      if numdiff:
+        return tf.test.compute_gradient(
+            [case.bias_op], [case.bias.shape], actual_op,
+            case.expected_output_shape())[0]
+      else:
+        return sess.run(tf.gradients(actual_op, [case.bias_op]))[0]
 
-      err = tf.test.compute_gradient_error([self.bias_op],
-                                           [self.bias.shape], actual_op,
-                                           TPC.expected_output_shape())
-      self.assertLess(err, 1e-2)
+  def _backward_theta(self, use_gpu=False, dtype=np.float32, numdiff=True):
+    case.init_ops(dtype=dtype)
+    with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu) as sess:
+      actual_op = flex_convolution(case.features_op,
+                                   case.position_op, case.neighborhood_op,
+                                   case.theta_op, case.bias_op)
 
-  def _backward_theta(self, use_gpu=False):
-    self.init_ops()
-    with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu):
-      actual_op = flex_convolution(self.features_op,
-                                   self.position_op, self.neighborhood_op,
-                                   self.theta_op, self.bias_op)
+      if numdiff:
+        return tf.test.compute_gradient(
+            [case.theta_op], [case.theta.shape], actual_op,
+            case.expected_output_shape())[0]
+      else:
+        return sess.run(tf.gradients(actual_op, [case.theta_op]))[0]
 
-      graph_theta_grad, num_theta_grad = tf.test.compute_gradient(
-          [self.theta_op], [self.theta.shape], actual_op,
-          TPC.expected_output_shape())[0]
-      summary(num_theta_grad, graph_theta_grad, 'self.theta')
+  def test_backward_features_cpu_float64(self):
+    actual, expected = self._backward_features(use_gpu=False, dtype=np.float64)
+    self.assertAllClose(actual, expected)
 
-      err = tf.test.compute_gradient_error([self.theta_op],
-                                           [self.theta.shape], actual_op,
-                                           TPC.expected_output_shape())
-      self.assertLess(err, 1e-2)
+  def test_backward_bias_cpu_float64(self):
+    actual, expected = self._backward_bias(use_gpu=False, dtype=np.float64)
+    self.assertAllClose(actual, expected)
+
+  def test_backward_theta_cpu_float64(self):
+    actual, expected = self._backward_theta(use_gpu=False, dtype=np.float64)
+    self.assertAllClose(actual, expected)
+
+  def test_backward_features_gpu_float64(self):
+    actual, expected = self._backward_features(use_gpu=True, dtype=np.float64)
+    self.assertAllClose(actual, expected)
+
+  def test_backward_bias_gpu_float64(self):
+    actual, expected = self._backward_bias(use_gpu=True, dtype=np.float64)
+    self.assertAllClose(actual, expected)
+
+  def test_backward_theta_gpu_float64(self):
+    actual, expected = self._backward_theta(use_gpu=True, dtype=np.float64)
+    self.assertAllClose(actual, expected)
+
+  def test_backward_features_gpu_float32(self, dtype=np.float32):
+    cpu = self._backward_features(use_gpu=False, dtype=dtype, numdiff=False)
+    gpu = self._backward_features(use_gpu=True, dtype=dtype, numdiff=False)
+    self.assertAllClose(cpu, gpu)
+
+  def test_backward_bias_gpu_float32(self, dtype=np.float32):
+    cpu = self._backward_bias(use_gpu=False, dtype=dtype, numdiff=False)
+    gpu = self._backward_bias(use_gpu=True, dtype=dtype, numdiff=False)
+    self.assertAllClose(cpu, gpu, 1e-5)
+
+  def test_backward_theta_gpu_float32(self, dtype=np.float32):
+    cpu = self._backward_theta(use_gpu=False, dtype=dtype, numdiff=False)
+    gpu = self._backward_theta(use_gpu=True, dtype=dtype, numdiff=False)
+    self.assertAllClose(cpu, gpu)
 
-  def test_backward_features(self):
-    self._backward_features(use_gpu=False)
-    self._backward_features(use_gpu=True)
 
-  def test_backward_bias(self):
-    self._backward_bias(use_gpu=False)
-    self._backward_bias(use_gpu=True)
 
-  def test_backward_theta(self):
-    self._backward_theta(use_gpu=False)
-    self._backward_theta(use_gpu=True)
 
 
 if __name__ == '__main__':
diff --git a/user_ops/test_flex_convolution_transpose.py b/user_ops/test_flex_convolution_transpose.py
index d627793..4dab694 100644
--- a/user_ops/test_flex_convolution_transpose.py
+++ b/user_ops/test_flex_convolution_transpose.py
@@ -18,98 +18,115 @@
 # Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
 
 
-from PointTestCase import TPC, PointTestCase, summary
+from misc import FakePointCloud, VerboseTestCase
 import tensorflow as tf
+import numpy as np
 from __init__ import flex_convolution_transpose
 
+case = FakePointCloud(B=2, N=32, K=4, Din=2, Dout=6, Dp=3)
 
-class FlexDeconvTest(PointTestCase):
+
+class FlexConvTest(VerboseTestCase):
   def __init__(self, methodName="runTest"):
-    super(FlexDeconvTest, self).__init__(methodName)
+    super(FlexConvTest, self).__init__(methodName)
 
   def _forward(self, use_gpu=False, force_gpu=False):
-    self.init_ops()
+    case.init_ops()
     with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu) as sess:
-      actual = flex_convolution_transpose(self.features_op,
-                                          self.position_op, self.neighborhood_op,
-                                          self.theta_op, self.bias_op)
-
-      actual = sess.run(actual)
-
+      actual_op = flex_convolution_transpose(case.features_op,
+                                             case.position_op,
+                                             case.neighborhood_op,
+                                             case.theta_op, case.bias_op)
+      actual = sess.run(actual_op)
     return actual
 
   def test_forward(self):
     cpu = self._forward(use_gpu=False)
     gpu = self._forward(use_gpu=True)
-    summary(cpu, gpu, 'self.features')
-    self.assertAllClose(cpu, gpu, rtol=1e-05)
+    self.assertAllClose(cpu, gpu, 1e-4)
 
-  def _backward_features(self, use_gpu=False):
-    self.init_ops()
-    with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu):
-      actual_op = flex_convolution_transpose(self.features_op,
-                                             self.position_op,
-                                             self.neighborhood_op,
-                                             self.theta_op, self.bias_op)
+  def _backward_features(self, use_gpu=False, dtype=np.float32, numdiff=True):
+    case.init_ops(dtype=dtype)
+    with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu) as sess:
+      actual_op = flex_convolution_transpose(case.features_op,
+                                             case.position_op,
+                                             case.neighborhood_op,
+                                             case.theta_op, case.bias_op)
+      if numdiff:
+        return tf.test.compute_gradient(
+            [case.features_op], [case.features.shape], actual_op,
+            case.expected_output_shape())[0]
+      else:
+        return sess.run(tf.gradients(actual_op, [case.features_op]))[0]
 
-      graph_features_grad, num_features_grad = tf.test.compute_gradient(
-          [self.features_op], [self.features.shape], actual_op,
-          TPC.expected_output_shape())[0]
-      summary(num_features_grad, graph_features_grad,
-              'self.features', max_outputs=1200000)
+  def _backward_bias(self, use_gpu=False, dtype=np.float32, numdiff=True):
+    case.init_ops(dtype=dtype)
+    with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu) as sess:
+      actual_op = flex_convolution_transpose(case.features_op,
+                                             case.position_op,
+                                             case.neighborhood_op,
+                                             case.theta_op, case.bias_op)
 
-      err = tf.test.compute_gradient_error([self.features_op],
-                                           [self.features.shape],
-                                           actual_op, TPC.expected_output_shape())
-      self.assertLess(err, 1e-2)
+      if numdiff:
+        return tf.test.compute_gradient(
+            [case.bias_op], [case.bias.shape], actual_op,
+            case.expected_output_shape())[0]
+      else:
+        return sess.run(tf.gradients(actual_op, [case.bias_op]))[0]
 
-  def _backward_bias(self, use_gpu=False):
-    self.init_ops()
-    with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu):
-      actual_op = flex_convolution_transpose(self.features_op,
-                                             self.position_op,
-                                             self.neighborhood_op,
-                                             self.theta_op, self.bias_op)
+  def _backward_theta(self, use_gpu=False, dtype=np.float32, numdiff=True):
+    case.init_ops(dtype=dtype)
+    with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu) as sess:
+      actual_op = flex_convolution_transpose(case.features_op,
+                                             case.position_op,
+                                             case.neighborhood_op,
+                                             case.theta_op, case.bias_op)
 
-      graph_bias_grad, num_bias_grad = tf.test.compute_gradient(
-          [self.bias_op], [self.bias.shape], actual_op,
-          TPC.expected_output_shape())[0]
-      summary(num_bias_grad, graph_bias_grad, 'self.bias')
+      if numdiff:
+        return tf.test.compute_gradient(
+            [case.theta_op], [case.theta.shape], actual_op,
+            case.expected_output_shape())[0]
+      else:
+        return sess.run(tf.gradients(actual_op, [case.theta_op]))[0]
 
-      err = tf.test.compute_gradient_error([self.bias_op],
-                                           [self.bias.shape], actual_op,
-                                           TPC.expected_output_shape())
-      self.assertLess(err, 1e-2)
+  def test_backward_features_cpu_float64(self):
+    actual, expected = self._backward_features(use_gpu=False, dtype=np.float64)
+    self.assertAllClose(actual, expected)
 
-  def _backward_theta(self, use_gpu=False):
-    self.init_ops()
-    with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu):
-      actual_op = flex_convolution_transpose(self.features_op,
-                                             self.position_op,
-                                             self.neighborhood_op,
-                                             self.theta_op, self.bias_op)
+  def test_backward_bias_cpu_float64(self):
+    actual, expected = self._backward_bias(use_gpu=False, dtype=np.float64)
+    self.assertAllClose(actual, expected)
 
-      graph_theta_grad, num_theta_grad = tf.test.compute_gradient(
-          [self.theta_op], [self.theta.shape], actual_op,
-          TPC.expected_output_shape())[0]
-      summary(num_theta_grad, graph_theta_grad, 'self.theta')
+  def test_backward_theta_cpu_float64(self):
+    actual, expected = self._backward_theta(use_gpu=False, dtype=np.float64)
+    self.assertAllClose(actual, expected)
 
-      err = tf.test.compute_gradient_error([self.theta_op],
-                                           [self.theta.shape], actual_op,
-                                           TPC.expected_output_shape())
-      self.assertLess(err, 1e-2)
+  def test_backward_features_gpu_float64(self):
+    actual, expected = self._backward_features(use_gpu=True, dtype=np.float64)
+    self.assertAllClose(actual, expected)
 
-  def test_backward_features(self):
-    self._backward_features(use_gpu=False)
-    self._backward_features(use_gpu=True)
+  def test_backward_bias_gpu_float64(self):
+    actual, expected = self._backward_bias(use_gpu=True, dtype=np.float64)
+    self.assertAllClose(actual, expected)
 
-  def test_backward_bias(self):
-    self._backward_bias(use_gpu=False)
-    self._backward_bias(use_gpu=True)
+  def test_backward_theta_gpu_float64(self):
+    actual, expected = self._backward_theta(use_gpu=True, dtype=np.float64)
+    self.assertAllClose(actual, expected)
 
-  def test_backward_theta(self):
-    self._backward_theta(use_gpu=False)
-    self._backward_theta(use_gpu=True)
+  def test_backward_features_gpu_float32(self, dtype=np.float32):
+    cpu = self._backward_features(use_gpu=False, dtype=dtype, numdiff=False)
+    gpu = self._backward_features(use_gpu=True, dtype=dtype, numdiff=False)
+    self.assertAllClose(cpu, gpu)
+
+  def test_backward_bias_gpu_float32(self, dtype=np.float32):
+    cpu = self._backward_bias(use_gpu=False, dtype=dtype, numdiff=False)
+    gpu = self._backward_bias(use_gpu=True, dtype=dtype, numdiff=False)
+    self.assertAllClose(cpu, gpu, 1e-5)
+
+  def test_backward_theta_gpu_float32(self, dtype=np.float32):
+    cpu = self._backward_theta(use_gpu=False, dtype=dtype, numdiff=False)
+    gpu = self._backward_theta(use_gpu=True, dtype=dtype, numdiff=False)
+    self.assertAllClose(cpu, gpu)
 
 
 if __name__ == '__main__':
diff --git a/user_ops/test_flex_pooling.py b/user_ops/test_flex_pooling.py
index c97a9bd..75507b6 100644
--- a/user_ops/test_flex_pooling.py
+++ b/user_ops/test_flex_pooling.py
@@ -18,70 +18,85 @@
 # Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
 
 
-from PointTestCase import PointTestCase
+from misc import FakePointCloud, VerboseTestCase
 import tensorflow as tf
 import numpy as np
 
 from __init__ import flex_pooling
 
+case = FakePointCloud(B=2, N=32, K=4, Din=2, Dout=6, Dp=3)
 
-class FlexPoolTest(PointTestCase):
-    def __init__(self, methodName="runTest"):
-        super(FlexPoolTest, self).__init__(methodName)
 
-    def _forward(self, use_gpu=False):
-        self.init_ops()
-        with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu) as sess:
-            actual_op, winner_op = flex_pooling(self.features_op, self.neighborhood_op)
-            actual = sess.run(actual_op)
-        return actual
+class FlexPoolTest(VerboseTestCase):
+  def __init__(self, methodName="runTest"):
+    super(FlexPoolTest, self).__init__(methodName)
 
-    def test_forward(self):
-        cpu = self._forward(use_gpu=False)
-        gpu = self._forward(use_gpu=True)
-        self.assertAllClose(cpu, gpu)
+  def _forward(self, use_gpu=False, dtype=np.float32):
+    case.init_ops(dtype=dtype)
+    with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu) as sess:
+      actual_op, winner_op = flex_pooling(
+          case.features_op, case.neighborhood_op)
+      actual = sess.run(actual_op)
+    return actual
 
-    def test_backward(self):
-        cpu, winner_cpu = self._backward(use_gpu=False)
-        gpu, winner_gpu = self._backward(use_gpu=True)
-        self.assertAllClose(cpu, gpu)
-        self.assertAllClose(cpu[winner_cpu == 0].sum(), 0)
-        self.assertAllClose(gpu[winner_gpu == 0].sum(), 0)
+  def test_forward_same_float32(self):
+    cpu = self._forward(use_gpu=False, dtype=np.float32)
+    gpu = self._forward(use_gpu=True, dtype=np.float32)
+    self.assertAllClose(cpu, gpu)
 
-    def _backward(self, use_gpu=False):
-        self.init_ops()
-        with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu) as sess:
-            actual_op, winner_op = flex_pooling(
-                self.features_op,
-                self.neighborhood_op)
+  def test_forward_same_float64(self):
+    cpu = self._forward(use_gpu=False, dtype=np.float64)
+    gpu = self._forward(use_gpu=True, dtype=np.float64)
+    self.assertAllClose(cpu, gpu)
 
-            graph_features_grad = tf.gradients(actual_op, [self.features_op])[0]
+  def test_backward_same_float32(self):
+    cpu, winner_cpu = self._backward(use_gpu=False, dtype=np.float32)
+    gpu, winner_gpu = self._backward(use_gpu=True, dtype=np.float32)
+    self.assertAllClose(winner_cpu, winner_gpu)
+    self.assertAllClose(cpu, gpu)
 
-            dx, winner = sess.run([graph_features_grad, winner_op])
-            return dx, winner
+  def test_backward_same_float64(self):
+    cpu, winner_cpu = self._backward(use_gpu=False, dtype=np.float64)
+    gpu, winner_gpu = self._backward(use_gpu=True, dtype=np.float64)
+    self.assertAllClose(winner_cpu, winner_gpu)
+    self.assertAllClose(cpu, gpu)
 
-    def _simple_backward(self, use_gpu=False):
-        # BN
-        x = np.array([[[1], [2], [5], [3]]]).transpose(0, 2, 1)
-        n = np.array([[[0, 1, 2, 3], [1, 2, 3, 0], [2, 3, 0, 1], [3, 0, 1, 2, ]]]).transpose(0, 2, 1)
+  def _backward(self, use_gpu=False, dtype=np.float32):
+    case.init_ops(dtype=dtype)
+    with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu) as sess:
+      actual_op, winner_op = flex_pooling(
+          case.features_op,
+          case.neighborhood_op)
 
-        x = tf.convert_to_tensor(x.astype(np.float32))
-        n = tf.convert_to_tensor(n.astype(np.int32))
+      graph_features_grad = tf.gradients(actual_op, [case.features_op])[0]
 
-        with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu) as sess:
-            actual_op, winner_op = flex_pooling(x, n)
-            graph_features_grad = tf.gradients(actual_op, [x])[0]
-            return sess.run(graph_features_grad)
+      dx, winner = sess.run([graph_features_grad, winner_op])
+      return dx, winner
 
-    def test_backward_simple(self):
-        cpu = self._simple_backward(use_gpu=False)
-        cpu[0, 0, 2] -= 4
-        self.assertEqual(cpu.sum(), 0)
+  def _simple_backward(self, use_gpu=False):
+    # BN
+    x = np.array([[[1], [2], [5], [3]]]).transpose(0, 2, 1)
+    n = np.array([[[0, 1, 2, 3], [1, 2, 3, 0], [2, 3, 0, 1],
+                   [3, 0, 1, 2, ]]]).transpose(0, 2, 1)
 
-        gpu = self._simple_backward(use_gpu=True)
-        gpu[0, 0, 2] -= 4
-        self.assertEqual(gpu.sum(), 0)
+    x = tf.convert_to_tensor(x.astype(np.float32))
+    n = tf.convert_to_tensor(n.astype(np.int32))
+
+    with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu) as sess:
+      actual_op, winner_op = flex_pooling(x, n)
+      graph_features_grad = tf.gradients(actual_op, [x])[0]
+      return sess.run(graph_features_grad)
+
+  def test_backward_simple_cpu(self):
+    cpu = self._simple_backward(use_gpu=False)
+    cpu[0, 0, 2] -= 4
+    self.assertEqual(cpu.sum(), 0)
+
+  def test_backward_simple_gpu(self):
+    gpu = self._simple_backward(use_gpu=True)
+    gpu[0, 0, 2] -= 4
+    self.assertEqual(gpu.sum(), 0)
 
 
 if __name__ == '__main__':
-    tf.test.main()
+  tf.test.main()
diff --git a/user_ops/test_knn_bruteforce.py b/user_ops/test_knn_bruteforce.py
new file mode 100644
index 0000000..5b5919a
--- /dev/null
+++ b/user_ops/test_knn_bruteforce.py
@@ -0,0 +1,66 @@
+#!/usr/bin/env python
+# -*- coding: utf-8 -*-
+
+# Copyright 2017 ComputerGraphics Tuebingen. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+# Authors: Fabian Groh, Patrick Wieschollek, Hendrik P.A. Lensch
+
+
+from misc import FakePointCloud, VerboseTestCase
+import tensorflow as tf
+import numpy as np
+from scipy.spatial.distance import pdist, squareform
+
+from __init__ import knn_bruteforce
+
+case = FakePointCloud(B=2, N=32, K=4, Din=2, Dout=6, Dp=3)
+case = FakePointCloud(B=1, N=4, K=2, Din=1, Dout=1, Dp=3)
+
+
+def python_bruteforce(positions, K):
+  # B, Dpos, N
+  all_neighbors = []
+  all_distances = []
+  for batch in positions:
+    distances = squareform(pdist(batch.T, 'euclidean'))
+    all_neighbors.append(np.argsort(distances, axis=1)[:, :K])
+    all_distances.append(np.sort(distances, axis=1)[:, :K])
+  return np.array(all_neighbors), np.array(all_distances)
+
+
+class KnnBruteforceTest(VerboseTestCase):
+  def __init__(self, methodName="runTest"):
+    super(KnnBruteforceTest, self).__init__(methodName)
+
+  def _forward(self, use_gpu=False):
+    case.init_ops()
+    expected_nn, expected_dist = python_bruteforce(case.position, K=4)
+
+    with self.test_session(use_gpu=use_gpu, force_gpu=use_gpu) as sess:
+      actual_nn, actual_dist, _ = knn_bruteforce(case.position_op, K=4)
+      actual_nn, actual_dist = sess.run([actual_nn, actual_dist])
+
+    self.assertAllClose(expected_dist, actual_dist)
+    self.assertAllClose(expected_nn, actual_nn)
+
+  def test_forward_cpu(self):
+    self._forward(use_gpu=False)
+
+  def test_forward_gpu(self):
+    self._forward(use_gpu=True)
+
+
+if __name__ == '__main__':
+  tf.test.main()