Add support for half

[cmpute]

Should fix #126.
I didn't test it yet since my model has other bugs using half precision...

[rusty1s]

Hi, are there any updates on this PR? The last time I was working with fp16, standard operators like + and * needed adoption, too.

[cmpute]

In this PR all functions that need to be adapted in pytorch-scatter has been fixed. Remaining functions are all pytorch built-in functions. I'll try to make pull request to pytorch and as long as those functions are adapted for half precision, this PR should work.

[rusty1s]

Cool! Please keep me updated:)

[cmpute]

I created a PR in pytorch. Actually I found that the tests only fail using CPU. When using CUDA, those functions missing in CPU are already implemented.

[murnanedaniel]

@cmpute, is this PR basically ready to merge? I'm very excited for this functionality to be incorporated! Let me know if there are any remaining jobs I can help with.

[cmpute]

@cmpute, is this PR basically ready to merge? I'm very excited for this functionality to be incorporated! Let me know if there are any remaining jobs I can help with.

I created a PR in pytorch, but it still need some time to be merged. I just need to write some test for that PR. After it's merged I think we're done~

[rusty1s]

Can you link to the PR?

[cmpute]

link

I guess Github already links it here under your last comment? pytorch/pytorch#39788

[speedcell4]

I tried this with Pytorch 1.8.1 and it works fine. Does this mean fp16 is fully supported by torch-scatter now?

from torch_scatter import scatter_add

import torch

a = torch.randn((5, 7), dtype=torch.half).cuda()
index = torch.randint(0, 10, (5,)).cuda()
print(scatter_add(a, index=index, dim=0))

[rusty1s]

Should be only supported for scatter_add and scatter_mean, not for scatter_max and scatter_mul yet.

[rusty1s]

Thank you very much. I fixed the remaining PyTorch compatibility issues and everything should work now.

[HelloWorldLTY]

Hi, I wonder if there are any example codes we can refer for amp training. Thanks a lot.

[rusty1s]

This might be helpful: https://github.com/pyg-team/pytorch_geometric/blob/master/benchmark/inference/inference_benchmark.py#L38

[jykr]

Hi, thanks for this update- I'm having trouble using half precision with HeteroLinear with following error (in using HGTConv). Am I doing something wrong, or does torch.ops.pyg not support half precision? Thank you.

  File "/data/pinello/PROJECTS/2022_12_GCPA/dynot/dynot/encoders.py", line 205, in forward
    _x_dict = conv(_x_dict, edge_dict)
              ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/pinello/SHARED_SOFTWARE/anaconda_latest/envs/jy_gcpa_torch4/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/pinello/SHARED_SOFTWARE/anaconda_latest/envs/jy_gcpa_torch4/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/pinello/PROJECTS/2022_12_GCPA/software/pytorch_geometric/torch_geometric/nn/conv/hgt_conv.py", line 194, in forward
    k, v, src_offset = self._construct_src_node_feat(
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/pinello/PROJECTS/2022_12_GCPA/software/pytorch_geometric/torch_geometric/nn/conv/hgt_conv.py", line 153, in _construct_src_node_feat
    k = self.k_rel(ks, type_vec).view(H, -1, D).transpose(0, 1)
        ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/pinello/SHARED_SOFTWARE/anaconda_latest/envs/jy_gcpa_torch4/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1553, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/pinello/SHARED_SOFTWARE/anaconda_latest/envs/jy_gcpa_torch4/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1562, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/pinello/PROJECTS/2022_12_GCPA/software/pytorch_geometric/torch_geometric/nn/dense/linear.py", line 342, in forward
    self._update_timing_cache(x, type_ptr, key)
  File "/data/pinello/SHARED_SOFTWARE/anaconda_latest/envs/jy_gcpa_torch4/lib/python3.12/site-packages/torch/utils/_contextlib.py", line 116, in decorate_context
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/data/pinello/PROJECTS/2022_12_GCPA/software/pytorch_geometric/torch_geometric/nn/dense/linear.py", line 301, in _update_timing_cache
    _ = self.forward_segmm(x, type_ptr)
        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/pinello/PROJECTS/2022_12_GCPA/software/pytorch_geometric/torch_geometric/nn/dense/linear.py", line 285, in forward_segmm
    return pyg_lib.ops.segment_matmul(x, type_ptr, self.weight)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/pinello/SHARED_SOFTWARE/anaconda_latest/envs/jy_gcpa_torch4/lib/python3.12/site-packages/pyg_lib/ops/__init__.py", line 174, in segment_matmul
    out = torch.ops.pyg.segment_matmul(inputs, ptr, other)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/pinello/SHARED_SOFTWARE/anaconda_latest/envs/jy_gcpa_torch4/lib/python3.12/site-packages/torch/_ops.py", line 1061, in __call__
    return self_._op(*args, **(kwargs or {}))
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: expected scalar type Float but found Half