Create `cugraph-equivariant` package #4036

tingyu66 · 2023-12-04T15:51:10Z

Bring up cugraph-equivariant package and add TensorProduct conv layers.

tingyu66 · 2024-01-19T14:51:05Z

@stadlmax I've updated the wrapper based on the new fused_tp kernels. Can you take a look? Tests won't pass until we have the new kernels reflected in pylibcugraphops nightly.

mariogeiger

There is a problem if mlp_fast_first_layer=False but src_scalars and dst_scalars are provided.

(optional) It would be nice if we can remove that mlp_fast_first_layer and simply do the thing when src_scalars and dst_scalars are provided.

python/cugraph-equivariant/cugraph_equivariant/nn/tensor_product_conv.py

mariogeiger · 2024-01-22T16:47:56Z

python/cugraph-equivariant/cugraph_equivariant/nn/tensor_product_conv.py

+        batch_norm: bool = True,
+        mlp_channels: Optional[Sequence[int]] = None,
+        mlp_activation: Optional[Callable[..., nn.Module]] = nn.GELU,
+        mlp_fast_first_layer: bool = False,


Why does this argument need to be here? Can't the special just happen when src_scalars and dst_scalars are provided?

I agree with Mario here. There are a few branches when it comes to computing the tp weights:

user direct input the precomputed weights to forward

feed only the edge_emb to the mlp to compute the weight. This could include the case that the node embeddings are already concatenated into the edge_emb, which we can't distinguish

As in point 2 but additionally have one node embedding tensor, i.e., the graph is non-directed

As in point 2 but additionally have separate src and dst node embedding tensors

Edit: we already handled cases 1, 2 and 4, despite some argument against the need for mlp_fast_first_layer in the init(). Do we want to handle case 3 explicitly in the API? i.e., if only src_scalars but not dst_scalars is given, do we want to assume the dst embedding are also index from src_scalars or do we always want the user to supply the same tensor to both src_scalars=node_attr, dst_scalars=node_attr? I guess it can't hurt to do the latter so it's up to you

I've updated the code to support arbitrary number of scalar features in [src, dst]_scalars. They can also be None if needed.

Regarding use case 3, users should input src_scalars=node_attr, dst_scalars=node_attr.

DejunL · 2024-01-22T16:52:12Z

python/cugraph-equivariant/cugraph_equivariant/nn/__init__.py

+
+from .tensor_product_conv import FullyConnectedTensorProductConv
+
+DiffDockTensorProductConv = FullyConnectedTensorProductConv


Why do we need this alias?

This is from one of our discussions on Monday, but it's totally optional. @mariogeiger Do you think we need the alias?

No need for me, the name FullyConnectedTensorProductConv is perfect by its own

python/cugraph-equivariant/cugraph_equivariant/nn/tensor_product_conv.py

DejunL · 2024-01-22T17:12:31Z

python/cugraph-equivariant/cugraph_equivariant/nn/tensor_product_conv.py

+        \sum_{b \in \mathcal{N}_a} Y\left(\hat{r}_{a b}\right)
+        \otimes_{\psi_{a b}} \mathbf{h}_b
+
+    where the path weights :math:`\psi_{a b}` are from user input to the forward()


In the Diffdock paper, 1) \psi is itself the MLP for computing the weights from the edge and node embeddings. 2) This implementation has the option to either directly input the weights to the forward() or compute the weights using the MLP (\psi) from the edge and/or the node embeddings

\psi denotes the weights while \Psi is the MLP. I have added a brief note to show the two options here

python/cugraph-equivariant/cugraph_equivariant/nn/tensor_product_conv.py

DejunL · 2024-01-22T17:35:25Z

python/cugraph-equivariant/cugraph_equivariant/nn/tensor_product_conv.py

+            mlp = []
+            for i in range(len(dims) - 1):
+                mlp.append(nn.Linear(dims[i], dims[i + 1]))
+                if mlp_activation is not None and i != len(dims) - 2:


When would mlp_activation be None?

python/cugraph-equivariant/cugraph_equivariant/nn/tensor_product_conv.py

tingyu66 · 2024-01-22T18:16:20Z

@mariogeiger @DejunL
I referred to the diffdock repo when implementing it and thought edge, src and dst sharing the same num_scalers would the only use case . If we do need support for various num_scalers for different component, I can make the change accordingly.

I will also remove the mlp_fast_first_layer flag, since it is only used to validate the dimensionality.

python/cugraph-equivariant/cugraph_equivariant/nn/tensor_product_conv.py

DejunL · 2024-01-22T19:30:19Z

python/cugraph-equivariant/cugraph_equivariant/nn/tensor_product_conv.py

+        in_irreps: o3.Irreps,
+        sh_irreps: o3.Irreps,
+        out_irreps: o3.Irreps,
+        batch_norm: bool = True,


In the Diffdock training code, we have customized batch_norm function. Instead of being a boolean option, can this take a callable that defaults to e3nn.nn.BatchNorm?

Sure. Will make the change.

I haven't come up with a neat way to support this because if we change the argument to be a callable (with the default value =Batchnorm), the customized batch_norm function must have the same function signature as e3nn.nn.Batchnorm. With that, I would suggest applying customized function outside of the conv layer.

Both the e3nn.nn.BatchNorm and the modified version we used are nn.module so I should have suggested change the type of the batch_norm argument to nn.module. But then it becomes a moot point since we can apply the BatchNorm outside of the conv layer like you said. Let me double check with Guoqing, who authored the modified version of BatchNorm

Just confirmed with Guoqing that the current API is OK. So it's up to you if you want to keep it as it is

Wouldn't it make more sense to just not have BN in the API then?

DejunL · 2024-01-22T19:47:58Z

@mariogeiger @DejunL I referred to the diffdock repo when implementing it and thought edge, src and dst sharing the same num_scalers would the only use case . If we do need support for various num_scalers for different component, I can make the change accordingly.

I will also remove the mlp_fast_first_layer flag, since it is only used to validate the dimensionality.

Yes, for now in the Diffdock model, edge_emb and src_scalars are really scalars. But in general, we don't want to constrain them to be scalar or of same dimensionality. Thank you!

rlratzel · 2024-01-24T05:18:16Z

ci/test_python.sh

+      --channel pytorch \
+      --channel nvidia \
+      cugraph-equivariant
+    pip install e3nn==0.5.1


I'm not sure, but would this be better off as a dependency in the py_test_cugraph_equivariant section of the dependencies.yaml file?

Ideally yes, but e3nn depends on PyTorch. Having that in pyproject.toml might pull wrong versions of pytorch for users. cugraph-dgl and -pyg's pyproject.toml does not have pytorch either (I guess for the same reason).

@tingyu66 do you want me to remove pytorch from the dependencies of e3nn?

python/cugraph-equivariant/cugraph_equivariant/tests/test_tensor_product_conv.py

…or_product_conv.py Co-authored-by: Mario Geiger <[email protected]>

…ct_conv.py Co-authored-by: Mario Geiger <[email protected]>

stadlmax

besides the minor comment w.r.t. having BN in the API, no further comments

mariogeiger

Same, looks good except that BN that I'm not super super happy about

BradReesWork · 2024-01-29T15:16:47Z

/merge

move scripts over to cugraph repo

a859eee

tingyu66 added this to the 24.02 milestone Dec 4, 2023

tingyu66 self-assigned this Dec 4, 2023

Merge branch 'branch-24.02' into cugraph-equivariant

b35b307

github-actions bot added the python label Dec 12, 2023

tingyu66 added 2 commits December 12, 2023 13:54

Merge branch 'branch-24.02' into cugraph-equivariant

38d8429

test conda build

d2186ee

github-actions bot added ci conda labels Dec 12, 2023

tingyu66 added 9 commits December 12, 2023 15:48

separate conda build into a new job

148cb4e

remove cpp channel

15404b5

fix workflow

3c58d7c

build wheel

834c1ba

merge in branch-24.02

d27a42b

Merge branch 'branch-24.02' into cugraph-equivariant

43b97b8

year

1f4077a

revise mlp support

6462cc3

test with mlps

2259a8d

tingyu66 added improvement Improvement / enhancement to an existing function non-breaking Non-breaking change labels Jan 12, 2024

tingyu66 added 5 commits January 12, 2024 17:21

test wheel

c525bef

Merge branch 'branch-24.02' into cugraph-equivariant

78fc221

absorb conda build into build_python script

3fd56dc

update based on cugraph-ops pr

924e55b

style

0a9c4b8

tingyu66 requested a review from stadlmax January 19, 2024 14:47

tingyu66 requested review from mariogeiger and DejunL January 22, 2024 16:27

mariogeiger suggested changes Jan 22, 2024

View reviewed changes

DejunL reviewed Jan 22, 2024

View reviewed changes

DejunL suggested changes Jan 22, 2024

View reviewed changes

DejunL reviewed Jan 22, 2024

View reviewed changes

python/cugraph-equivariant/cugraph_equivariant/nn/tensor_product_conv.py Show resolved Hide resolved

DejunL reviewed Jan 22, 2024

View reviewed changes

python/cugraph-equivariant/cugraph_equivariant/nn/tensor_product_conv.py Outdated Show resolved Hide resolved

DejunL reviewed Jan 22, 2024

View reviewed changes

tingyu66 added 2 commits January 22, 2024 21:15

Merge branch 'branch-24.02' into cugraph-equivariant

9547bd5

address comments

b9ce1ca

tingyu66 marked this pull request as ready for review January 23, 2024 19:32

tingyu66 requested review from a team as code owners January 23, 2024 19:32

tingyu66 changed the title ~~[DRAFT] Create cugraph-equivariant package~~ Create cugraph-equivariant package Jan 23, 2024

tingyu66 added 4 commits January 23, 2024 14:34

mlp_activation accept sequences, imp docstrings

4457eff

clean up

65ceca3

update workflow files

f848852

correct test script

c0abf8c

rlratzel reviewed Jan 24, 2024

View reviewed changes

mariogeiger reviewed Jan 24, 2024

View reviewed changes

python/cugraph-equivariant/cugraph_equivariant/tests/test_tensor_product_conv.py Outdated Show resolved Hide resolved

tingyu66 and others added 3 commits January 24, 2024 11:19

Update python/cugraph-equivariant/cugraph_equivariant/tests/test_tens…

40b972c

…or_product_conv.py Co-authored-by: Mario Geiger <[email protected]>

Update python/cugraph-equivariant/cugraph_equivariant/nn/tensor_produ…

6dff47a

…ct_conv.py Co-authored-by: Mario Geiger <[email protected]>

format

3ee809b

AyodeAwe approved these changes Jan 24, 2024

View reviewed changes

DejunL approved these changes Jan 25, 2024

View reviewed changes

tingyu66 requested a review from a team January 26, 2024 17:27

stadlmax approved these changes Jan 29, 2024

View reviewed changes

mariogeiger approved these changes Jan 29, 2024

View reviewed changes

rlratzel approved these changes Jan 29, 2024

View reviewed changes

rapids-bot bot merged commit 3ff2abd into rapidsai:branch-24.02 Jan 29, 2024
109 checks passed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Create `cugraph-equivariant` package #4036

Create `cugraph-equivariant` package #4036

tingyu66 commented Dec 4, 2023

tingyu66 commented Jan 19, 2024

mariogeiger left a comment

mariogeiger Jan 22, 2024

DejunL Jan 22, 2024 •

edited

Loading

tingyu66 Jan 24, 2024

DejunL Jan 22, 2024

tingyu66 Jan 23, 2024

mariogeiger Jan 29, 2024

DejunL Jan 22, 2024

tingyu66 Jan 23, 2024

DejunL Jan 22, 2024

tingyu66 commented Jan 22, 2024

DejunL Jan 22, 2024

tingyu66 Jan 22, 2024

tingyu66 Jan 24, 2024 •

edited

Loading

DejunL Jan 25, 2024

DejunL Jan 25, 2024

stadlmax Jan 29, 2024

DejunL commented Jan 22, 2024

rlratzel Jan 24, 2024

tingyu66 Jan 24, 2024

mariogeiger Jan 29, 2024

stadlmax left a comment

mariogeiger left a comment

BradReesWork commented Jan 29, 2024


		from .tensor_product_conv import FullyConnectedTensorProductConv

		DiffDockTensorProductConv = FullyConnectedTensorProductConv

Create cugraph-equivariant package #4036

Create cugraph-equivariant package #4036

Conversation

tingyu66 commented Dec 4, 2023

tingyu66 commented Jan 19, 2024

mariogeiger left a comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

DejunL Jan 22, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tingyu66 commented Jan 22, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

tingyu66 Jan 24, 2024 • edited Loading

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

DejunL commented Jan 22, 2024

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

stadlmax left a comment

Choose a reason for hiding this comment

mariogeiger left a comment

Choose a reason for hiding this comment

BradReesWork commented Jan 29, 2024

Create `cugraph-equivariant` package #4036

Create `cugraph-equivariant` package #4036

DejunL Jan 22, 2024 •

edited

Loading

tingyu66 Jan 24, 2024 •

edited

Loading