diff options
author | Vitaly Fedyunin <vitalyf@fb.com> | 2019-04-22 17:48:43 -0700 |
---|---|---|
committer | Facebook Github Bot <facebook-github-bot@users.noreply.github.com> | 2019-04-22 17:52:47 -0700 |
commit | 8be6d5ffd89854a4ef4296544e1042a453396cfd (patch) | |
tree | fb67763f26390415d39cc9a8b7ec33fea19d3d3e /torch | |
parent | 5a796d15be92a58a2387e0c05d9f0f613893df74 (diff) | |
download | pytorch-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.py | 7 |
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' |