29 using namespace tensorflow;
33 const Tensor& radius =
context->input(1);
34 OP_REQUIRES(
context, TensorShapeUtils::IsScalar(radius.shape()),
35 errors::InvalidArgument(
"radius must be scalar, got shape ",
36 radius.shape().DebugString()));
38 const Tensor& points_row_splits =
context->input(2);
40 const Tensor& hash_table_size_factor_tensor =
context->input(3);
43 TensorShapeUtils::IsScalar(
44 hash_table_size_factor_tensor.shape()),
45 errors::InvalidArgument(
46 "hash_table_size_factor must be scalar, got shape ",
47 hash_table_size_factor_tensor.shape().DebugString()));
48 const double hash_table_size_factor =
49 hash_table_size_factor_tensor.scalar<
double>()();
51 Dim num_points(
"num_points");
52 Dim batch_size(
"batch_size");
56 std::vector<uint32_t> hash_table_splits(batch_size.
value() + 1, 0);
57 for (
int i = 0; i < batch_size.
value(); ++i) {
58 int64_t num_points_i = points_row_splits.flat<int64>()(i + 1) -
59 points_row_splits.flat<int64>()(i);
61 int64_t hash_table_size = std::min<int64_t>(
62 std::max<int64_t>(hash_table_size_factor * num_points_i, 1),
64 hash_table_splits[i + 1] = hash_table_splits[i] + hash_table_size;
67 Tensor* hash_table_index = 0;
68 TensorShape hash_table_index_shape({num_points.
value()});
70 context->allocate_output(0, hash_table_index_shape,
73 Tensor* hash_table_cell_splits = 0;
74 TensorShape hash_table_cell_splits_shape(
75 {hash_table_splits.back() + 1});
77 context->allocate_output(1, hash_table_cell_splits_shape,
78 &hash_table_cell_splits));
80 Tensor* out_hash_table_splits = 0;
81 TensorShape out_hash_table_splits_shape({batch_size.
value() + 1});
83 context->allocate_output(2, out_hash_table_splits_shape,
84 &out_hash_table_splits));
85 for (
size_t i = 0; i < hash_table_splits.size(); ++i) {
86 out_hash_table_splits->flat<uint32_t>()(i) = hash_table_splits[i];
90 *hash_table_index, *hash_table_cell_splits);