29#include "../TensorFlowHelper.h"
32#include "tensorflow/core/framework/op.h"
33#include "tensorflow/core/framework/op_kernel.h"
34#include "tensorflow/core/lib/core/errors.h"
38namespace fixed_radius_search_opkernel {
41template <
class T,
class TIndex>
42class OutputAllocator {
46 void AllocIndices(TIndex** ptr,
size_t num) {
47 using namespace tensorflow;
50 TensorShape shape({int64_t(num)});
51 OP_REQUIRES_OK(
context,
context->allocate_output(0, shape, &tensor));
53 auto flat_tensor = tensor->flat<TIndex>();
54 *ptr = (TIndex*)flat_tensor.data();
57 void AllocDistances(T** ptr,
size_t num) {
58 using namespace tensorflow;
61 TensorShape shape({int64_t(num)});
62 OP_REQUIRES_OK(
context,
context->allocate_output(2, shape, &tensor));
63 auto flat_tensor = tensor->flat<T>();
64 *ptr = flat_tensor.data();
68 tensorflow::OpKernelContext*
context;
71class FixedRadiusSearchOpKernel :
public tensorflow::OpKernel {
73 explicit FixedRadiusSearchOpKernel(
74 tensorflow::OpKernelConstruction* construction)
75 : OpKernel(construction) {
76 using namespace tensorflow;
79 std::string metric_str;
80 OP_REQUIRES_OK(construction,
81 construction->GetAttr(
"metric", &metric_str));
82 if (metric_str ==
"L1")
84 else if (metric_str ==
"L2")
89 OP_REQUIRES_OK(construction,
90 construction->GetAttr(
"ignore_query_point",
91 &ignore_query_point));
93 OP_REQUIRES_OK(construction, construction->GetAttr(
"return_distances",
97 void Compute(tensorflow::OpKernelContext*
context)
override {
98 using namespace tensorflow;
99 static_assert(
sizeof(int64) ==
sizeof(int64_t),
100 "int64 type is not compatible");
106 OP_REQUIRES(
context, TensorShapeUtils::IsScalar(radius.shape()),
107 errors::InvalidArgument(
"radius must be scalar, got shape ",
108 radius.shape().DebugString()));
120 Dim num_points(
"num_points");
121 Dim num_queries(
"num_queries");
122 Dim batch_size(
"batch_size");
123 Dim num_cells(
"num_cells");
132 Tensor* query_neighbors_row_splits = 0;
133 TensorShape query_neighbors_row_splits_shape(
134 {queries.shape().dim_size(0) + 1});
136 1, query_neighbors_row_splits_shape,
137 &query_neighbors_row_splits));
140 queries_row_splits, hash_table_splits, hash_table_index,
141 hash_table_cell_splits, *query_neighbors_row_splits);
144 virtual void Kernel(tensorflow::OpKernelContext*
context,
145 const tensorflow::Tensor&
points,
146 const tensorflow::Tensor& queries,
147 const tensorflow::Tensor& radius,
148 const tensorflow::Tensor& points_row_splits,
149 const tensorflow::Tensor& queries_row_splits,
150 const tensorflow::Tensor& hash_table_splits,
151 const tensorflow::Tensor& hash_table_index,
152 const tensorflow::Tensor& hash_table_cell_splits,
153 tensorflow::Tensor& query_neighbors_row_splits) = 0;
157 bool ignore_query_point;
158 bool return_distances;
#define CHECK_SHAPE(tensor,...)
Definition: TorchHelper.h:205
ImGuiContext * context
Definition: Window.cpp:95
Class for dimensions for which the value should be inferred.
Definition: ShapeChecking.h:69
Definition: FixedRadiusIndex.cpp:35
Metric
Supported metrics.
Definition: NeighborSearchCommon.h:38
@ Linf
Definition: NeighborSearchCommon.h:38
@ L1
Definition: NeighborSearchCommon.h:38
@ L2
Definition: NeighborSearchCommon.h:38
Definition: ShapeChecking.h:35