Home | History | Annotate | Download | only in nearest_neighbor 
      1 # Description:
      2 #   Tensorflow ops for nearest neighbor queries etc.
      3 
      4 package(default_visibility = ["//tensorflow:__subpackages__"])
      5 
      6 licenses(["notice"])  # Apache 2.0
      7 
      8 exports_files(["LICENSE"])
      9 
     10 load("//tensorflow:tensorflow.bzl", "tf_custom_op_py_library")
     11 load(
     12     "//tensorflow:tensorflow.bzl",
     13     "tf_cc_test",
     14     "tf_custom_op_library",
     15     "tf_gen_op_libs",
     16     "tf_gen_op_wrapper_py",
     17     "tf_kernel_library",
     18     "tf_py_test",
     19 )
     20 
     21 tf_custom_op_library(
     22     name = "python/ops/_nearest_neighbor_ops.so",
     23     srcs = [
     24         "kernels/hyperplane_lsh_probes.cc",
     25         "ops/nearest_neighbor_ops.cc",
     26     ],
     27     deps = [
     28         ":hyperplane_lsh_probes",
     29     ],
     30 )
     31 
     32 tf_gen_op_libs(
     33     op_lib_names = ["nearest_neighbor_ops"],
     34 )
     35 
     36 tf_gen_op_wrapper_py(
     37     name = "nearest_neighbor_ops_pywrapper",
     38     deps = ["nearest_neighbor_ops_op_lib"],
     39 )
     40 
     41 tf_custom_op_py_library(
     42     name = "nearest_neighbor_py",
     43     srcs = ["__init__.py"] + glob(["python/ops/*.py"]),
     44     dso = [":python/ops/_nearest_neighbor_ops.so"],
     45     kernels = [":nearest_neighbor_ops_kernels"],
     46     srcs_version = "PY2AND3",
     47     visibility = ["//visibility:public"],
     48     deps = [
     49         "//tensorflow/contrib/util:util_py",
     50         "//tensorflow/python:framework_for_generated_wrappers",
     51         "//tensorflow/python:platform",
     52     ],
     53 )
     54 
     55 tf_kernel_library(
     56     name = "nearest_neighbor_ops_kernels",
     57     srcs = ["kernels/hyperplane_lsh_probes.cc"],
     58     deps = [
     59         ":hyperplane_lsh_probes",
     60         ":nearest_neighbor_ops_op_lib",
     61         "//tensorflow/core:framework",
     62         "//tensorflow/core:lib",
     63         "//third_party/eigen3",
     64     ],
     65 )
     66 
     67 cc_library(
     68     name = "heap",
     69     hdrs = ["kernels/heap.h"],
     70 )
     71 
     72 tf_cc_test(
     73     name = "heap_test",
     74     size = "small",
     75     srcs = ["kernels/heap_test.cc"],
     76     deps = [
     77         ":heap",
     78         "//tensorflow/core:test_main",
     79         "//tensorflow/core/kernels:ops_testutil",
     80     ],
     81 )
     82 
     83 cc_library(
     84     name = "hyperplane_lsh_probes",
     85     hdrs = ["kernels/hyperplane_lsh_probes.h"],
     86     deps = [
     87         ":heap",
     88         "//third_party/eigen3",
     89     ],
     90 )
     91 
     92 tf_cc_test(
     93     name = "hyperplane_lsh_probes_test_cc",
     94     size = "small",
     95     srcs = ["kernels/hyperplane_lsh_probes_test.cc"],
     96     deps = [
     97         ":hyperplane_lsh_probes",
     98         "//tensorflow/core:test",
     99         "//tensorflow/core:test_main",
    100         "//tensorflow/core:testlib",
    101         "//tensorflow/core/kernels:ops_testutil",
    102     ],
    103 )
    104 
    105 tf_py_test(
    106     name = "hyperplane_lsh_probes_test",
    107     size = "small",
    108     srcs = ["python/kernel_tests/hyperplane_lsh_probes_test.py"],
    109     additional_deps = [
    110         ":nearest_neighbor_py",
    111         "//tensorflow/python:client_testlib",
    112     ],
    113 )
    114