diff options
Diffstat (limited to 'tensorflow/core/ops/array_ops.cc')
| -rw-r--r-- | tensorflow/core/ops/array_ops.cc | 17 |
1 files changed, 17 insertions, 0 deletions
diff --git a/tensorflow/core/ops/array_ops.cc b/tensorflow/core/ops/array_ops.cc index a1027f1422..39b92464cb 100644 --- a/tensorflow/core/ops/array_ops.cc +++ b/tensorflow/core/ops/array_ops.cc @@ -1203,6 +1203,23 @@ REGISTER_OP("UniqueWithCounts") return Status::OK(); }); +REGISTER_OP("UniqueWithCountsV2") + .Input("x: T") + .Input("axis: Taxis") + .Output("y: T") + .Output("idx: out_idx") + .Output("count: out_idx") + .Attr("T: type") + .Attr("Taxis: {int32,int64} = DT_INT64") + .Attr("out_idx: {int32, int64} = DT_INT32") + .SetShapeFn([](InferenceContext* c) { + auto uniq = c->Vector(InferenceContext::kUnknownDim); + c->set_output(0, uniq); + c->set_output(1, c->input(0)); + c->set_output(2, uniq); + return Status::OK(); + }); + namespace { Status ShapeShapeFn(InferenceContext* c) { |