mirror of
https://github.com/wassname/Flex-Convolution.git
synced 2026-09-09 11:14:10 +08:00
42 lines
1.3 KiB
C++
42 lines
1.3 KiB
C++
// ComputerGraphics Tuebingen, 2018
|
|
|
|
#include "tensorflow/core/framework/op.h"
|
|
#include "tensorflow/core/framework/shape_inference.h"
|
|
|
|
namespace tensorflow {
|
|
|
|
using ::tensorflow::shape_inference::InferenceContext;
|
|
using ::tensorflow::shape_inference::ShapeHandle;
|
|
|
|
REGISTER_OP("KnnBruteforce")
|
|
.Input("position: T") // position: each datapoint in nd space [B,Dp, N].
|
|
.Output("neighborhood_out: NBtype") // neighborhood_out: all K nearest
|
|
// neighbors [B, N, K].
|
|
.Output("distances: T") // distances: all K nearest distances
|
|
// [B, N, K].
|
|
.Output("timings: T") // timings:
|
|
// [1]
|
|
.Attr("K: int")
|
|
.Attr("return_timings: bool = false")
|
|
.Attr("T: realnumbertype")
|
|
.Attr("NBtype: {int32} = DT_INT32")
|
|
.SetShapeFn([](::tensorflow::shape_inference::InferenceContext* c) {
|
|
const auto position = c->input(0);
|
|
const auto neighborhood_out = c->input(1);
|
|
|
|
int K;
|
|
c->GetAttr("K", &K);
|
|
|
|
auto B = c->Dim(position, 0);
|
|
auto N = c->Dim(position, 2);
|
|
|
|
c->set_output(0, c->MakeShape({B, N, K}));
|
|
c->set_output(1, c->MakeShape({B, N, K}));
|
|
|
|
return Status::OK();
|
|
});
|
|
|
|
} // namespace tensorflow
|
|
|
|
// doc: K: number of neighbors.
|