1 | // Copyright 2018 The TensorFlow Authors. All Rights Reserved. |
2 | // |
3 | // Licensed under the Apache License, Version 2.0 (the "License"); you may not |
4 | // use this file except in compliance with the License. You may obtain a copy |
5 | // 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, WITHOUT |
11 | // WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the |
12 | // License for the specific language governing permissions and limitations under |
13 | // the License. |
14 | // ============================================================================== |
15 | |
16 | #include "tensorflow/core/framework/common_shape_fns.h" |
17 | #include "tensorflow/core/framework/op.h" |
18 | |
19 | namespace tensorflow { |
20 | |
21 | REGISTER_OP("KmeansPlusPlusInitialization" ) |
22 | .Input("points: float32" ) |
23 | .Input("num_to_sample: int64" ) |
24 | .Input("seed: int64" ) |
25 | .Input("num_retries_per_sample: int64" ) |
26 | .Output("samples: float32" ) |
27 | .SetShapeFn(shape_inference::UnknownShape); |
28 | |
29 | REGISTER_OP("KMC2ChainInitialization" ) |
30 | .Input("distances: float32" ) |
31 | .Input("seed: int64" ) |
32 | .Output("index: int64" ) |
33 | .SetShapeFn(shape_inference::ScalarShape); |
34 | |
35 | REGISTER_OP("NearestNeighbors" ) |
36 | .Input("points: float32" ) |
37 | .Input("centers: float32" ) |
38 | .Input("k: int64" ) |
39 | .Output("nearest_center_indices: int64" ) |
40 | .Output("nearest_center_distances: float32" ) |
41 | .SetShapeFn(shape_inference::UnknownShape); |
42 | |
43 | } // namespace tensorflow |
44 | |