🚀 Feature

Add torch.unique support to torch.jit.script so that it is possible to jit-script graph pooling in using pytorch geometric.

This was mentioned back in #12206.

@suo

Motivation

I am writing models using pytorch_geometric that I would like to jit-script. I have started by making automatic GNN model synthesis that makes jittable convolution operations, but we also need to do graph pooling to have the full range of functionality. This latter bit requires torch.unique in python to be jit-scriptable.

Pitch

I would like:

@torch.jit.script
def f(x):
    return torch.unique(x)

to not result in:

Traceback (most recent call last):
  File "quicktest.py", line 3, in <module>
    @torch.jit.script
  File "/data/gaoxiang/anaconda3/lib/python3.6/site-packages/torch/jit/__init__.py", line 639, in script
    graph = _jit_script_compile(ast, rcb)
RuntimeError: 
unknown builtin op:
@torch.jit.script
def f(x):
    return torch.unique(x)
           ~~~~~~~~~~~~ <--- HERE

but rather a working jit function!

Alternatives

Writing everything in C++ for torch_geometric, which is a significant time burden for developers, where this could be fixed easily and centrally.

Additional context

This is supporting an effort here:
pyg-team/pytorch_geometric#1191

to make pytorch_geometric easily jittable without major rewrites or changes of user code.

cc @suo