1# Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2# 3# Licensed under the Apache License, Version 2.0 (the "License"); 4# you may not use this file except in compliance with the License. 5# You may obtain a copy of the License at 6# 7# http://www.apache.org/licenses/LICENSE-2.0 8# 9# Unless required by applicable law or agreed to in writing, software 10# distributed under the License is distributed on an "AS IS" BASIS, 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12# See the License for the specific language governing permissions and 13# limitations under the License. 14# ============================================================================== 15"""Wrappers for nearest neighbor operations.""" 16 17from __future__ import absolute_import 18from __future__ import division 19from __future__ import print_function 20 21from tensorflow.contrib.util import loader 22from tensorflow.python.framework import ops 23from tensorflow.python.platform import resource_loader 24 25_nearest_neighbor_ops = loader.load_op_library( 26 resource_loader.get_path_to_datafile("_nearest_neighbor_ops.so")) 27 28 29def hyperplane_lsh_probes(point_hyperplane_product, 30 num_tables, 31 num_hyperplanes_per_table, 32 num_probes, 33 name=None): 34 """Computes probes for the hyperplane hash. 35 36 The op supports multiprobing, i.e., the number of requested probes can be 37 larger than the number of tables. In that case, the same table can be probed 38 multiple times. 39 40 The first `num_tables` probes are always the primary hashes for each table. 41 42 Args: 43 point_hyperplane_product: a matrix of inner products between the hyperplanes 44 and the points to be hashed. These values should not be quantized so that 45 we can correctly compute the probing sequence. The expected shape is 46 `batch_size` times `num_tables * num_hyperplanes_per_table`, i.e., each 47 element of the batch corresponds to one row of the matrix. 48 num_tables: the number of tables to compute probes for. 49 num_hyperplanes_per_table: the number of hyperplanes per table. 50 num_probes: the requested number of probes per table. 51 name: A name prefix for the returned tensors (optional). 52 53 Returns: 54 probes: the output matrix of probes. Size `batch_size` times `num_probes`. 55 table_ids: the output matrix of tables ids. Size `batch_size` times 56 `num_probes`. 57 """ 58 return _nearest_neighbor_ops.hyperplane_lsh_probes(point_hyperplane_product, 59 num_tables, 60 num_hyperplanes_per_table, 61 num_probes, 62 name=name) 63 64ops.NotDifferentiable("HyperplaneLSHProbes") 65