summaryrefslogtreecommitdiff
path: root/torch
diff options
context:
space:
mode:
authorVitaly Fedyunin <vitalyf@fb.com>2019-04-22 17:48:43 -0700
committerFacebook Github Bot <facebook-github-bot@users.noreply.github.com>2019-04-22 17:52:47 -0700
commit8be6d5ffd89854a4ef4296544e1042a453396cfd (patch)
treefb67763f26390415d39cc9a8b7ec33fea19d3d3e /torch
parent5a796d15be92a58a2387e0c05d9f0f613893df74 (diff)
downloadpytorch-8be6d5ffd89854a4ef4296544e1042a453396cfd.tar.gz
pytorch-8be6d5ffd89854a4ef4296544e1042a453396cfd.tar.bz2
pytorch-8be6d5ffd89854a4ef4296544e1042a453396cfd.zip
Add onnx support for _unique2 operator
Summary: Pull Request resolved: https://github.com/pytorch/pytorch/pull/19582 Reviewed By: ezyang, jamesr66a Differential Revision: D15037375 fbshipit-source-id: 6060476925bf02fa07f852054e06d2107f046e38
Diffstat (limited to 'torch')
-rw-r--r--torch/onnx/symbolic.py7
1 files changed, 7 insertions, 0 deletions
diff --git a/torch/onnx/symbolic.py b/torch/onnx/symbolic.py
index ba06e981f5..8f87391646 100644
--- a/torch/onnx/symbolic.py
+++ b/torch/onnx/symbolic.py
@@ -1247,6 +1247,13 @@ def _unique(g, input, sorted, return_inverse):
return_inverse_i=return_inverse, outputs=2)
+@parse_args('v', 'i', 'i', 'i')
+def _unique2(g, input, sorted, return_inverse, return_counts):
+ return g.op("ATen", input, operator_s="_unique2", sorted_i=sorted,
+ return_inverse_i=return_inverse, return_counts_i=return_counts,
+ outputs=3)
+
+
# Metaprogram symbolics for each ATen native specialized cast operator.
# For e.g. we specify a function named `_cast_uint8_t` that instantiates an
# ONNX cast node with `to` attribute 'UINT8'