25 #include "../TensorFlowHelper.h" 27 #include "tensorflow/core/framework/op.h" 28 #include "tensorflow/core/framework/op_kernel.h" 29 #include "tensorflow/core/lib/core/errors.h" 35 class RadiusSearchOpKernel :
public tensorflow::OpKernel {
37 explicit RadiusSearchOpKernel(
38 tensorflow::OpKernelConstruction* construction)
39 : OpKernel(construction) {
42 std::string metric_str;
43 OP_REQUIRES_OK(construction,
44 construction->GetAttr(
"metric", &metric_str));
45 if (metric_str ==
"L1")
50 OP_REQUIRES_OK(construction,
51 construction->GetAttr(
"ignore_query_point",
52 &ignore_query_point));
54 OP_REQUIRES_OK(construction, construction->GetAttr(
"return_distances",
56 OP_REQUIRES_OK(construction,
57 construction->GetAttr(
"normalize_distances",
58 &normalize_distances));
61 void Compute(tensorflow::OpKernelContext* context)
override {
63 static_assert(
sizeof(int64) ==
sizeof(int64_t),
64 "int64 type is not compatible");
66 const Tensor&
points = context->input(0);
67 const Tensor& queries = context->input(1);
68 const Tensor& radii = context->input(2);
69 const Tensor& points_row_splits = context->input(3);
70 const Tensor& queries_row_splits = context->input(4);
74 Dim num_points(
"num_points");
75 Dim num_queries(
"num_queries");
76 Dim batch_size(
"batch_size");
80 CHECK_SHAPE(context, points_row_splits, batch_size + 1);
81 CHECK_SHAPE(context, queries_row_splits, batch_size + 1);
84 Tensor* query_neighbors_row_splits = 0;
85 TensorShape query_neighbors_row_splits_shape(
86 {queries.shape().dim_size(0) + 1});
87 OP_REQUIRES_OK(context, context->allocate_output(
88 1, query_neighbors_row_splits_shape,
89 &query_neighbors_row_splits));
91 Kernel(context, points, queries, radii, points_row_splits,
92 queries_row_splits, *query_neighbors_row_splits);
95 virtual void Kernel(tensorflow::OpKernelContext* context,
96 const tensorflow::Tensor& points,
97 const tensorflow::Tensor& queries,
98 const tensorflow::Tensor& radius,
99 const tensorflow::Tensor& points_row_splits,
100 const tensorflow::Tensor& queries_row_splits,
101 tensorflow::Tensor& query_neighbors_row_splits) = 0;
105 bool ignore_query_point;
106 bool return_distances;
107 bool normalize_distances;
#define CHECK_SHAPE(tensor,...)
Definition: TorchHelper.h:204
Metric
Supported metrics.
Definition: NeighborSearchCommon.h:38
Definition: NanoFlannIndex.h:55
Class for dimensions for which the value should be inferred.
Definition: ShapeChecking.h:69
Definition: ContinuousConv.h:35
int points
Definition: FilePCD.cpp:73
Definition: NanoFlannIndex.h:55
Definition: ShapeChecking.h:35