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 16 #include "tensorflow/core/framework/op.h" 17 18 namespace tensorflow { 19 20 REGISTER_OP("HyperplaneLSHProbes") 21 .Attr("CoordinateType: {float, double}") 22 .Input("point_hyperplane_product: CoordinateType") 23 .Input("num_tables: int32") 24 .Input("num_hyperplanes_per_table: int32") 25 .Input("num_probes: int32") 26 .Output("probes: int32") 27 .Output("table_ids: int32") 28 .Doc(R"doc( 29 Computes probes for the hyperplane hash. 30 31 The op supports multiprobing, i.e., the number of requested probes can be 32 larger than the number of tables. In that case, the same table can be probed 33 multiple times. 34 35 The first `num_tables` probes are always the primary hashes for each table. 36 37 point_hyperplane_product: a matrix of inner products between the hyperplanes 38 and the points to be hashed. These values should not be quantized so that we 39 can correctly compute the probing sequence. The expected shape is 40 `batch_size` times `num_tables * num_hyperplanes_per_table`, i.e., each 41 element of the batch corresponds to one row of the matrix. 42 num_tables: the number of tables to compute probes for. 43 num_hyperplanes_per_table: the number of hyperplanes per table. 44 num_probes: the requested number of probes per table. 45 probes: the output matrix of probes. Size `batch_size` times `num_probes`. 46 table_ids: the output matrix of tables ids. Size `batch_size` times `num_probes`. 47 )doc"); 48 49 } // namespace tensorflow 50