diff options
| author |  A. Unique TensorFlower <gardener@tensorflow.org> | 2018-09-20 12:37:30 -0700 |
|---|---|---|
| committer | TensorFlower Gardener <gardener@tensorflow.org> | 2018-09-20 12:44:24 -0700 |
| commit | 4aa639c0cbb47f4707f735e0cc80f4c39506d928 (patch) | |
| tree | ee87e1b72e0c6a9482e1005f88706f7c193e7a33 /tensorflow/core/ops | |
| parent | 350effcc2fd95c723c92267cf13fcd38777a2a98 (diff) | |
Add searchsorted (ie lower/upper bound) op.
PiperOrigin-RevId: 213863392
Diffstat (limited to 'tensorflow/core/ops')
| -rw-r--r-- | tensorflow/core/ops/array_ops.cc | 28 |
1 files changed, 28 insertions, 0 deletions
diff --git a/tensorflow/core/ops/array_ops.cc b/tensorflow/core/ops/array_ops.cc index 7dbb18aa5d..c24950643f 100644 --- a/tensorflow/core/ops/array_ops.cc +++ b/tensorflow/core/ops/array_ops.cc @@ -2916,6 +2916,34 @@ Status ScatterNdShape(InferenceContext* c) { } // namespace +REGISTER_OP("UpperBound") + .Input("sorted_inputs: T") + .Input("values: T") + .Output("output: out_type") + .Attr("T: type") + .Attr("out_type: {int32, int64} = DT_INT32") + .SetShapeFn([](InferenceContext* c) { + ShapeHandle unused_shape; + TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 2, &unused_shape)); + TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 2, &unused_shape)); + c->set_output(0, c->input(1)); + return Status::OK(); + }); + +REGISTER_OP("LowerBound") + .Input("sorted_inputs: T") + .Input("values: T") + .Output("output: out_type") + .Attr("T: type") + .Attr("out_type: {int32, int64} = DT_INT32") + .SetShapeFn([](InferenceContext* c) { + ShapeHandle unused_shape; + TF_RETURN_IF_ERROR(c->WithRank(c->input(0), 2, &unused_shape)); + TF_RETURN_IF_ERROR(c->WithRank(c->input(1), 2, &unused_shape)); + c->set_output(0, c->input(1)); + return Status::OK(); + }); + REGISTER_OP("ScatterNd") .Input("indices: Tindices") .Input("updates: T") |