10#include "../TensorFlowHelper.h"
13#include "tensorflow/core/framework/op.h"
14#include "tensorflow/core/framework/op_kernel.h"
15#include "tensorflow/core/lib/core/errors.h"
19namespace fixed_radius_search_opkernel {
22template <
class T,
class TIndex>
23class OutputAllocator {
27 void AllocIndices(TIndex** ptr,
size_t num) {
28 using namespace tensorflow;
31 TensorShape shape({int64_t(num)});
32 OP_REQUIRES_OK(
context,
context->allocate_output(0, shape, &tensor));
34 auto flat_tensor = tensor->flat<TIndex>();
35 *ptr = (TIndex*)flat_tensor.data();
38 void AllocDistances(T** ptr,
size_t num) {
39 using namespace tensorflow;
42 TensorShape shape({int64_t(num)});
43 OP_REQUIRES_OK(
context,
context->allocate_output(2, shape, &tensor));
44 auto flat_tensor = tensor->flat<T>();
45 *ptr = flat_tensor.data();
49 tensorflow::OpKernelContext*
context;
52class FixedRadiusSearchOpKernel :
public tensorflow::OpKernel {
54 explicit FixedRadiusSearchOpKernel(
55 tensorflow::OpKernelConstruction* construction)
56 : OpKernel(construction) {
57 using namespace tensorflow;
60 std::string metric_str;
61 OP_REQUIRES_OK(construction,
62 construction->GetAttr(
"metric", &metric_str));
63 if (metric_str ==
"L1")
65 else if (metric_str ==
"L2")
70 OP_REQUIRES_OK(construction,
71 construction->GetAttr(
"ignore_query_point",
72 &ignore_query_point));
74 OP_REQUIRES_OK(construction, construction->GetAttr(
"return_distances",
78 void Compute(tensorflow::OpKernelContext*
context)
override {
79 using namespace tensorflow;
80 static_assert(
sizeof(int64) ==
sizeof(int64_t),
81 "int64 type is not compatible");
87 OP_REQUIRES(
context, TensorShapeUtils::IsScalar(radius.shape()),
88 errors::InvalidArgument(
"radius must be scalar, got shape ",
89 radius.shape().DebugString()));
101 Dim num_points(
"num_points");
102 Dim num_queries(
"num_queries");
103 Dim batch_size(
"batch_size");
104 Dim num_cells(
"num_cells");
113 Tensor* query_neighbors_row_splits = 0;
114 TensorShape query_neighbors_row_splits_shape(
115 {queries.shape().dim_size(0) + 1});
117 1, query_neighbors_row_splits_shape,
118 &query_neighbors_row_splits));
121 queries_row_splits, hash_table_splits, hash_table_index,
122 hash_table_cell_splits, *query_neighbors_row_splits);
125 virtual void Kernel(tensorflow::OpKernelContext*
context,
126 const tensorflow::Tensor&
points,
127 const tensorflow::Tensor& queries,
128 const tensorflow::Tensor& radius,
129 const tensorflow::Tensor& points_row_splits,
130 const tensorflow::Tensor& queries_row_splits,
131 const tensorflow::Tensor& hash_table_splits,
132 const tensorflow::Tensor& hash_table_index,
133 const tensorflow::Tensor& hash_table_cell_splits,
134 tensorflow::Tensor& query_neighbors_row_splits) = 0;
138 bool ignore_query_point;
139 bool return_distances;
#define CHECK_SHAPE(tensor,...)
Definition TorchHelper.h:186
ImGuiContext * context
Definition Window.cpp:76
Class for dimensions for which the value should be inferred.
Definition ShapeChecking.h:50
Definition FixedRadiusIndex.cpp:16
Metric
Supported metrics.
Definition NeighborSearchCommon.h:19
@ Linf
Definition NeighborSearchCommon.h:19
@ L1
Definition NeighborSearchCommon.h:19
@ L2
Definition NeighborSearchCommon.h:19
Definition ShapeChecking.h:16