ROCm support

[paulamartingonzalez]

Hi! Do you support ROCm AMD GPUs for this GNN package?

Thanks!

[janpfeifer]

hi @paulamartingonzalez ,

In principle GPUs supported by TensorFlow will be supported as well. But we haven't tried it yet.

The tricky part is that for the common GNN case, the library treat things as sparse and rely on tf.gather and the various unsorted segment operations:

https://github.com/tensorflow/gnn/blob/bcc256d47acc9d43f3af2cd4812f19cfde52607d/tensorflow_gnn/graph/graph_tensor_ops.py#L346

Assuming those TF ops are supported for ROCm AMD in TensorFlow, I would expect it to just work, but I'm not sure.

Let us know if you try it -- pls be patient, we are still working on the OSS version, it's very "alpha".

[arnoegw]

Catching up on old issues... It still is like @janpfeifer explained: broadcast and pool (the key operations to move data across GraphTensors) turn into tf.gather and tf.unsorted_segment_<op>, and are executed as usual by core TensorFlow.

At this time, we have no plans to go beyond TensorFlow's built-in support for GPU types, so I'm afraid there is no actionable feature request or defect report here. I hope we could answer your original question. To keep our list of issues manageable, please allow me mark this one as closed.