Repository: dmlc/dgl
Pull request: [Model] Add `dgl.nn.CuGraphRelGraphConv` model
## Description
<!-- Brief description. Refer to the related issues if existed.
It'll be great if relevant reviewers can be assigned as well.-->

This PR adds a RGCN model, `dgl.nn.CuGraphRelGraphConv`, that uses the accelerated sparse aggregation primitives in cugraph-ops. The model requires `pylibcugraphops>=22.12`.

## Checklist
Please feel free to remove inapplicable items for your PR.
- [x] The PR title starts with [$CATEGORY] (such as [NN], [Model], [Doc], [Feature]])
- [x] I've leverage the [tools](https://docs.google.com/document/d/1iHyj7zlmygKSk5gBPsqIqL5ASPzJSPREaNT_QdsiYA4/edit) to beautify the python and c++ code.
- [x] The PR is complete and small, read the [Google eng practice (CL equals to PR)](https://google.github.io/eng-practices/review/developer/small-cls.html) to understand more about small PR. In DGL, we consider PRs with less than 200 lines of core code change are small (example, test and documentation could be exempted).
- [x] All changes have test coverage
- [x] Code is well-documented
- [x] To the best of my knowledge, examples are either not affected by this change, or have been fixed to be compatible with this change
- [x] Related issue is referred in this PR
- [x] If the PR is for a new model/paper, I've updated the example index [here](../examples/README.md).

## Changes
- [x] New nn.Module: `dgl.nn.CuGraphRelGraphConv`
- [x] Add test that verifies the results against `dgl.nn.RelGraphConv`
- [x] Update API docs 

### Accuracy: Entity Classification

test accuracy after 100 epochs, using 20 bases

| Dataset   |  Mini-batch
| ------------ |  ------
| aifb         | ~0.92
| mutag     | ~0.58
| bgs         | ~0.70
| am          | ~0.64

## Performance
Hardware: 32-core AMD Ryzen Threadripper Pro 3975WX, NVIDIA RTX A6000
Benchmark: Mini-batch `entity_sample.py` example:
- batch_size: 100, graph idtype is int64.
- 2-layer RGCN, with fanouts: (4,4), hidden embedding dimension: 16, num_bases: 20
- The timings are for 100 epochs and the average across 3 identical runs.

|       | time_dgl (s) | time_cugraph-ops (s) |
|:-----:|:------------:|:--------------------:|
|  aifb |     4.24     |         2.61         |
| mutag |     5.65     |         3.68         |
|  bgs  |     4.19     |         2.70         |
|   am  |     19.11    |         13.71        |

During forward aggregation, most of the time (~70%) is spent on the conversion from DGL graph (in COO format) to cugraph-ops MFG type (in CSR format), indicated by `get_csr` and `create_mfg` entries in the following output. The actual forward aggregation function (agg_fwd_func) only takes 1/5 of the total time. 

```
Total Time (ns)  Instances   Avg (ns)     Med (ns)    Min (ns)  Max (ns)   StdDev (ns)   Style       Range    
---------------  ---------  -----------  -----------  --------  ---------  -----------  -------  -------------
    834,193,721        814  1,024,808.0  1,015,957.0   957,171  2,264,271     96,833.9  PushPop  total_agg_fwd
    395,471,358        814    485,837.0    474,488.0   453,493  1,577,778     78,678.5  PushPop  create_mfg   
    182,373,749        814    224,046.4    233,108.5   195,217    398,820     22,571.4  PushPop  get_csr      
    159,149,002        814    195,514.7    193,233.5   185,399    488,559     13,795.3  PushPop  agg_fwd_func
```
Hence, if DGL natively supports CSR representation of graphs, we can expect a much greater speed-up from using `dgl.nn.CuGraphRelGraphConv`.

### Notes

- Partially resolves #4669 
- This PR supersedes #4701
- Link to internal reviews https://github.com/rapidsai/dgl/pull/15.
- Closes https://github.com/rapidsai/cugraph-ops/issues/179

Problem statement:
[FEA] Create RGCN model using cugraph-ops aggregators
## 🚀 Feature
This integration creates an accelerated R-GCN model that uses the aggregation functions from `pylibcugraphops`, a Python wrapper around cugraph-ops aggregators.

Per our previous discussion, this new model `dgl.contrib.cugraph.nn.RelGraphConv` should mimic the interface of existing model `dgl.nn.RelGraphConv` and should first go to the `contrib` folder first. This accelerated model will only support basis decomposition (in terms of weight regularization) and requires DGL sampled graph as input. It is not intended for full-graph use cases.

## Motivation

<!-- Please outline the motivation for the proposal. Is your feature request
related to a problem? e.g., I'm always frustrated when [...]. If this is
related to another GitHub issue, please link here too -->

Provide a faster RGCN model on GPU to DGL users

Hints from issue discussion:
This issue has been automatically marked as stale due to lack of activity. It will be closed if no further activity occurs. Thank you

Commit message: update docs

Patch to review:
diff --git a/docs/source/api/python/nn-pytorch.rst b/docs/source/api/python/nn-pytorch.rst
index c550e5444c61..fffbe527f6f3 100644
--- a/docs/source/api/python/nn-pytorch.rst
+++ b/docs/source/api/python/nn-pytorch.rst
@@ -14,6 +14,7 @@ Conv Layers
     ~dgl.nn.pytorch.conv.GraphConv
     ~dgl.nn.pytorch.conv.EdgeWeightNorm
     ~dgl.nn.pytorch.conv.RelGraphConv
+    ~dgl.nn.pytorch.conv.CuGraphRelGraphConv
     ~dgl.nn.pytorch.conv.TAGConv
     ~dgl.nn.pytorch.conv.GATConv
     ~dgl.nn.pytorch.conv.GATv2Conv
diff --git a/python/dgl/nn/pytorch/conv/__init__.py b/python/dgl/nn/pytorch/conv/__init__.py
index f9d3b415a466..28386eb9813c 100644
--- a/python/dgl/nn/pytorch/conv/__init__.py
+++ b/python/dgl/nn/pytorch/conv/__init__.py
@@ -66,3 +66,10 @@
     "PNAConv",
     "DGNConv",
 ]
+
+try:
+    from .cugraph_relgraphconv import CuGraphRelGraphConv
+except ImportError:
+    pass
+else:
+    __all__.append('CuGraphRelGraphConv')
diff --git a/python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py b/python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py
new file mode 100644
index 000000000000..7ed4b0ef3d71
--- /dev/null
+++ b/python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py
@@ -0,0 +1,388 @@
+"""Torch Module for Relational graph convolution layer using the
+aggregation primitives in cugraph-ops"""
+# pylint: disable=no-member, arguments-differ, invalid-name, too-many-arguments
+import math
+
+import torch as th
+from torch import nn
+
+try:
+    from pylibcugraphops.legacy.aggregators.node_level import (
+        agg_hg_basis_post_bwd_int32,
+        agg_hg_basis_post_bwd_int64,
+        agg_hg_basis_post_fwd_int32,
+        agg_hg_basis_post_fwd_int64,
+    )
+    from pylibcugraphops.legacy.structure.graph_types import (
+        message_flow_graph_hg_csr_int32,
+        message_flow_graph_hg_csr_int64,
+    )
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "dgl.nn.CuGraphRelGraphConv requires pylibcugraphops to be installed."
+    )
+
+CUDA_SM_PER_BLOCK = 49152
+
+
+class RelGraphConvAgg(th.autograd.Function):
+    r"""Custom autograd function for R-GCN aggregation layer that uses the
+    aggregation functions in cugraph-ops."""
+
+    @staticmethod
+    def forward(ctx, g, fanout, num_rels, edge_types, feat, coeff):
+        r"""Compute the forward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        fanout : int
+            Maximum in-degree of nodes.
+        num_rels : int
+            Number of relations.
+        edge_types : torch.Tensor
+            A 1D tensor of edge types. Note that pylibcugraphops only accepts
+            edge_types in int32, so any input of other integer types will be casted
+            into int32, thus introducing some overhead. Pass in int32 tensor for
+            best performance.
+        feat : torch.Tensor, dtype=torch.float32, requires_grad=True
+            Tensor of source node features. Shape: (num_src_nodes, in_feat).
+        coeff : torch.Tensor, dtype=torch.float32, requires_grad=True
+            The coefficient matrix used in basis-decomposition regularization.
+            Shape: (num_rels, num_bases). It should be set to ``None``
+            when ``regularizer=None``.
+
+        Returns
+        -------
+        agg_output : torch.Tensor, dtype=torch.float32
+            Aggregation output. Shape: (num_dst_nodes, num_rels * in_feat)
+            when ``regularizer=None``; Shape: (num_dst_nodes, num_bases * in_feat)
+            when ``regularizer='basis'``.
+        """
+        ctx.device = g.device
+        if g.idtype == th.int32:
+            mfg_csr_func = message_flow_graph_hg_csr_int32
+            agg_fwd_func = agg_hg_basis_post_fwd_int32
+        elif g.idtype == th.int64:
+            mfg_csr_func = message_flow_graph_hg_csr_int64
+            agg_fwd_func = agg_hg_basis_post_fwd_int64
+        else:
+            raise TypeError(
+                f"Supported ID type: torch.int32 or torch.int64, "
+                f"but got {g.idtype}"
+            )
+        ctx.idtype = g.idtype
+
+        _in_feat = feat.shape[-1]
+        indptr, indices, edge_ids = g.adj_sparse("csc")
+        # edge_ids is in a mixed order, need to permutate incoming etypes
+        # and stash the result for backward propagation
+        ctx.edge_types_int32 = edge_types[edge_ids.long()].int()
+        # node_types are not needed by the post-variant RGCN aggregation functions
+        _num_node_types = 0
+        _out_node_types = _in_node_types = None
+
+        mfg = mfg_csr_func(
+            fanout,
+            g.dstnodes(),
+            g.srcnodes(),
+            indptr,
+            indices,
+            _num_node_types,
+            num_rels,
+            _out_node_types,
+            _in_node_types,
+            ctx.edge_types_int32,
+        )
+        ctx.mfg = mfg
+
+        if coeff is None:
+            leading_dimension = num_rels * _in_feat
+        else:
+            _num_bases = coeff.shape[-1]
+            leading_dimension = _num_bases * _in_feat
+
+        agg_output = th.empty(
+            g.num_dst_nodes(),
+            leading_dimension,
+            dtype=th.float32,
+            device=ctx.device,
+        )
+        if coeff is None:
+            agg_fwd_func(agg_output, feat.detach(), mfg)
+        else:
+            agg_fwd_func(
+                agg_output,
+                feat.detach(),
+                mfg,
+                weights_combination=coeff.detach()
+            )
+
+        ctx.save_for_backward(feat, coeff)
+        return agg_output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        r"""Compute the backward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        grad_output : torch.Tensor, dtype=torch.float32
+            Gradient of loss function w.r.t output.
+        """
+        feat, coeff = ctx.saved_tensors
+
+        if ctx.idtype == th.int32:
+            agg_bwd_func = agg_hg_basis_post_bwd_int32
+        else:
+            agg_bwd_func = agg_hg_basis_post_bwd_int64
+
+        grad_feat = th.empty_like(feat, dtype=th.float32, device=ctx.device)
+        if coeff is None:
+            grad_coeff = None
+            agg_bwd_func(grad_feat, grad_output, feat.detach(), ctx.mfg)
+        else:
+            grad_coeff = th.empty_like(
+                coeff, dtype=th.float32, device=ctx.device
+            )
+            agg_bwd_func(
+                grad_feat,
+                grad_output,
+                feat.detach(),
+                ctx.mfg,
+                output_weight_gradient=grad_coeff,
+                weights_combination=coeff.detach()
+            )
+
+        return None, None, None, None, grad_feat, grad_coeff
+
+
+class CuGraphRelGraphConv(nn.Module):
+    r"""An accelerated relational graph convolution layer using the
+    highly-optimized aggregation primitives in cugraph-ops.
+    See :class:`dgl.nn.pytorch.conv.RelGraphConv` for mathematical model.
+
+    .. note::
+        This is an **experimental** feature.
+        Compared with :class:`dgl.nn.pytorch.conv.RelGraphConv`, this model:
+
+        * Only works on cuda devices.
+        * Only support basis-decomposition regularization.
+        * Requires an extra argument `fanout` as input, since the current model is
+          designed for sampled-graph (message-flow-graph) use-cases.
+          Full-graph support will be added in upcoming releases.
+
+    Parameters
+    ----------
+    in_feat : int
+        Input feature size.
+    out_feat : int
+        Output feature size.
+    num_rels : int
+        Number of relations.
+    fanout : int
+        Maximum number of sampled neighbors of an destination node;
+        i.e, maximum in degree of destination nodes
+    regularizer : str, optional
+        Which weight regularizer to use ("basis" or ``None``):
+         - "basis" is for basis-decomposition.
+         - ``None`` applies no regularization.
+        Default: ``None``.
+    num_bases : int, optional
+        Number of bases. Only used when "basis" regularizer is applied.
+        Default: ``None``.
+    bias : bool, optional
+        True if bias is added. Default: ``True``.
+    activation : callable, optional
+        Activation function. Default: ``None``.
+    self_loop : bool, optional
+        True to include self loop message. Default: ``True``.
+    dropout : float, optional
+        Dropout rate. Default: ``0.0``
+    layer_norm: bool, optional
+        True to add layer norm. Default: ``False``
+
+    Examples
+    --------
+    >>> import dgl
+    >>> from dgl.dataloading import NeighborSampler, DataLoader
+    >>> from dgl.nn import CuGraphRelGraphConv
+    >>> import torch.nn.functional as F
+    ...
+    >>> device = 'cuda'
+    >>> fanouts = [5, 6]
+    >>> sampler = NeighborSampler(fanouts)
+    >>> dataloader = DataLoader(g, train_nid, sampler,
+    ...     device=device, batch_size=1024)
+    >>> conv1 = CuGraphRelGraphConv(in_dim, h_dim, num_rels, fanouts[0],
+    ...     regularizer='basis', num_bases=10).to(device)
+    >>> conv2 = CuGraphRelGraphConv(h_dim, out_dim, num_rels, fanouts[1],
+    ...     regularizer='basis', num_bases=10).to(device)
+    >>> for input_nodes, output_nodes, blocks in dataloader:
+    ...     h = conv1(blocks[0], x, blocks[0].edata[dgl.ETYPE])
+    ...     h = F.relu(h)
+    ...     h = conv2(blocks[1], h, blocks[1].edata[dgl.ETYPE])
+    """
+    def __init__(
+        self,
+        in_feat,
+        out_feat,
+        num_rels,
+        fanout,
+        regularizer=None,
+        num_bases=None,
+        bias=True,
+        activation=None,
+        self_loop=True,
+        dropout=0.0,
+        layer_norm=False,
+    ):
+        super().__init__()
+        self.in_feat = in_feat
+        self.out_feat = out_feat
+        self.num_rels = num_rels
+        self.fanout = fanout
+
+        # regularizer (see dgl.nn.pytorch.linear.TypedLinear)
+        if regularizer is None:
+            self.W = nn.Parameter(th.Tensor(num_rels, in_feat, out_feat))
+            self.coeff = None
+        elif regularizer == "basis":
+            if num_bases is None:
+                raise ValueError(
+                    'Missing "num_bases" for basis regularization.'
+                )
+            self.W = nn.Parameter(th.Tensor(num_bases, in_feat, out_feat))
+            self.coeff = nn.Parameter(th.Tensor(num_rels, num_bases))
+            self.num_bases = num_bases
+        else:
+            raise ValueError(
+                f"Supported regularizer options: 'basis', "
+                f"but got {regularizer}"
+            )
+        self.regularizer = regularizer
+        self.reset_parameters()
+
+        # others
+        self.bias = bias
+        self.activation = activation
+        self.self_loop = self_loop
+        self.layer_norm = layer_norm
+
+        # bias
+        if self.bias:
+            self.h_bias = nn.Parameter(th.Tensor(out_feat))
+            nn.init.zeros_(self.h_bias)
+
+        # layer norm
+        if self.layer_norm:
+            self.layer_norm_weight = nn.LayerNorm(
+                out_feat, elementwise_affine=True
+            )
+
+        # weight for self_loop
+        if self.self_loop:
+            self.loop_weight = nn.Parameter(th.Tensor(in_feat, out_feat))
+            nn.init.xavier_uniform_(
+                self.loop_weight, gain=nn.init.calculate_gain("relu")
+            )
+
+        self.dropout = nn.Dropout(dropout)
+
+    def reset_parameters(self):
+        with th.no_grad():
+            if self.regularizer is None:
+                nn.init.uniform_(
+                    self.W,
+                    -1 / math.sqrt(self.in_feat),
+                    1 / math.sqrt(self.in_feat),
+                )
+            elif self.regularizer == "basis":
+                nn.init.uniform_(
+                    self.W,
+                    -1 / math.sqrt(self.in_feat),
+                    1 / math.sqrt(self.in_feat),
+                )
+                nn.init.xavier_uniform_(
+                    self.coeff, gain=nn.init.calculate_gain("relu")
+                )
+            else:
+                raise ValueError(
+                    f"Supported regularizer options: 'basis', "
+                    f"but got {self.regularizer}"
+                )
+
+    def forward(self, g, feat, etypes, norm=None):
+        r"""Forward computation.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        feat : torch.Tensor
+            A 2D tensor of node features. Shape: :math:`(|V|, D_{in})`.
+        etypes : torch.Tensor or list[int]
+            An 1D integer tensor of edge types. Shape: :math:`(|E|,)`.
+        norm : torch.Tensor, optional
+            An 1D tensor of edge norm value.  Shape: :math:`(|E|,)`.
+
+        Returns
+        -------
+        torch.Tensor
+            New node features. Shape: :math:`(|V|, D_{out})`.
+        """
+        _device = next(self.parameters()).device
+        if _device.type != "cuda":
+            raise RuntimeError(
+                f"dgl.nn.CuGraphRelGraphConv requires "
+                f"the model on device 'cuda', but got '{_device.type}'"
+            )
+        if _device != g.device:
+            raise RuntimeError(
+                f"Expected model and graph on the same device, "
+                f"but got '{_device}' and '{g.device}'"
+            )
+        if _device != etypes.device:
+            raise RuntimeError(
+                f"Expected model and etypes on the same device, "
+                f"but got '{_device}' and '{etypes.device}'"
+            )
+        if _device != feat.device:
+            raise RuntimeError(
+                f"Expected model and feature tensor on the same device, "
+                f"but got '{_device}' and '{feat.device}'"
+            )
+        # check shared memory per block size
+        _sm_size = (
+            self.fanout * 2 + 3 + self.num_rels * self.num_bases
+            if self.regularizer == "basis"
+            else 0
+        )
+        _sm_size *= 8 if g.idtype == th.int64 else 4
+        if _sm_size > CUDA_SM_PER_BLOCK:
+            raise MemoryError(
+                f"Failed to allocate {_sm_size} bytes shared memory on CUDA, "
+                f"larger than the limit: {CUDA_SM_PER_BLOCK} bytes. "
+                f"Try reducing fanout or num_bases."
+            )
+        with g.local_scope():
+            g.srcdata["h"] = feat
+            if norm is not None:
+                g.edata["norm"] = norm
+            # message passing
+            h = RelGraphConvAgg.apply(
+                g, self.fanout, self.num_rels, etypes, feat, self.coeff
+            )
+            h = h @ self.W.view(-1, self.out_feat)
+            # apply bias and activation
+            if self.layer_norm:
+                h = self.layer_norm_weight(h)
+            if self.bias:
+                h = h + self.h_bias
+            if self.self_loop:
+                h = h + feat[: g.num_dst_nodes()] @ self.loop_weight
+            if self.activation:
+                h = self.activation(h)
+            h = self.dropout(h)
+            return h
diff --git a/tests/cugraph/cugraph-ops/test_cugraph_relgraphconv.py b/tests/cugraph/cugraph-ops/test_cugraph_relgraphconv.py
new file mode 100644
index 000000000000..7ab1fac52c9f
--- /dev/null
+++ b/tests/cugraph/cugraph-ops/test_cugraph_relgraphconv.py
@@ -0,0 +1,90 @@
+import pytest
+import torch
+import dgl
+from dgl.nn import CuGraphRelGraphConv
+from dgl.nn import RelGraphConv
+
+regularizers = [None, "basis"]
+device = "cuda"
+
+
+def generate_graph():
+    u = torch.tensor([0, 1, 0, 2, 3, 0, 4, 0, 5, 0, 6, 7, 0, 8, 9])
+    v = torch.tensor([1, 9, 2, 9, 9, 4, 9, 5, 9, 6, 9, 9, 8, 9, 0])
+    g = dgl.graph((u, v))
+    num_rels = 3
+    g.edata[dgl.ETYPE] = torch.randint(num_rels, (g.num_edges(),))
+    return g
+
+
+@pytest.mark.parametrize("regularizer", regularizers)
+def test_full_graph(regularizer):
+    in_feat, out_feat, num_rels, num_bases = 10, 2, 3, 2
+    kwargs = {
+        "num_bases": num_bases,
+        "regularizer": regularizer,
+        "bias": False,
+        "self_loop": False,
+    }
+    g = generate_graph().to(device)
+    fanout = g.in_degrees().max().item()
+    feat = torch.ones(g.num_nodes(), in_feat).to(device)
+
+    torch.manual_seed(0)
+    conv1 = RelGraphConv(in_feat, out_feat, num_rels, **kwargs).to(device)
+
+    torch.manual_seed(0)
+    conv2 = CuGraphRelGraphConv(
+        in_feat, out_feat, num_rels, fanout, **kwargs
+    ).to(device)
+
+    out1 = conv1(g, feat, g.edata[dgl.ETYPE])
+    out2 = conv2(g, feat, g.edata[dgl.ETYPE])
+
+    assert torch.allclose(out1, out2, atol=1e-06)
+
+    grad_out = torch.rand_like(out1)
+    out1.backward(grad_out)
+    out2.backward(grad_out)
+    assert torch.allclose(conv1.linear_r.W.grad, conv2.W.grad, atol=1e-6)
+    if regularizer is not None:
+        assert torch.allclose(
+            conv1.linear_r.coeff.grad, conv2.coeff.grad, atol=1e-6
+        )
+
+
+@pytest.mark.parametrize("regularizer", regularizers)
+def test_mfg(regularizer):
+    in_feat, out_feat, num_rels, num_bases = 10, 2, 3, 2
+    kwargs = {
+        "num_bases": num_bases,
+        "regularizer": regularizer,
+        "bias": False,
+        "self_loop": False,
+    }
+    g = generate_graph().to(device)
+    block = dgl.to_block(g)
+    fanout = block.in_degrees().max().item()
+    feat = torch.ones(g.num_nodes(), in_feat).to(device)
+
+    torch.manual_seed(0)
+    conv1 = RelGraphConv(in_feat, out_feat, num_rels, **kwargs).to(device)
+
+    torch.manual_seed(0)
+    conv2 = CuGraphRelGraphConv(
+        in_feat, out_feat, num_rels, fanout, **kwargs
+    ).to(device)
+
+    out1 = conv1(block, feat[block.srcdata[dgl.NID]], block.edata[dgl.ETYPE])
+    out2 = conv2(block, feat[block.srcdata[dgl.NID]], block.edata[dgl.ETYPE])
+
+    assert torch.allclose(out1, out2, atol=1e-06)
+
+    grad_out = torch.rand_like(out1)
+    out1.backward(grad_out)
+    out2.backward(grad_out)
+    assert torch.allclose(conv1.linear_r.W.grad, conv2.W.grad, atol=1e-6)
+    if regularizer is not None:
+        assert torch.allclose(
+            conv1.linear_r.coeff.grad, conv2.coeff.grad, atol=1e-6
+        )


Code review comments:
diff_hunk: @@ -0,0 +1,388 @@
+"""Torch Module for Relational graph convolution layer using the
+aggregation primitives in cugraph-ops"""
+# pylint: disable=no-member, arguments-differ, invalid-name, too-many-arguments
+import math
+
+import torch as th
+from torch import nn
+
+try:
+    from pylibcugraphops.legacy.aggregators.node_level import (
+        agg_hg_basis_post_bwd_int32,
+        agg_hg_basis_post_bwd_int64,
+        agg_hg_basis_post_fwd_int32,
+        agg_hg_basis_post_fwd_int64,
+    )
+    from pylibcugraphops.legacy.structure.graph_types import (
+        message_flow_graph_hg_csr_int32,
+        message_flow_graph_hg_csr_int64,
+    )
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "dgl.nn.CuGraphRelGraphConv requires pylibcugraphops to be installed."
+    )
+
+CUDA_SM_PER_BLOCK = 49152
+
+
+class RelGraphConvAgg(th.autograd.Function):
+    r"""Custom autograd function for R-GCN aggregation layer that uses the
+    aggregation functions in cugraph-ops."""
+
+    @staticmethod
+    def forward(ctx, g, fanout, num_rels, edge_types, feat, coeff):
+        r"""Compute the forward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        fanout : int
+            Maximum in-degree of nodes.
+        num_rels : int
+            Number of relations.
+        edge_types : torch.Tensor
+            A 1D tensor of edge types. Note that pylibcugraphops only accepts
+            edge_types in int32, so any input of other integer types will be casted
+            into int32, thus introducing some overhead. Pass in int32 tensor for
+            best performance.
+        feat : torch.Tensor, dtype=torch.float32, requires_grad=True
+            Tensor of source node features. Shape: (num_src_nodes, in_feat).
+        coeff : torch.Tensor, dtype=torch.float32, requires_grad=True
+            The coefficient matrix used in basis-decomposition regularization.
+            Shape: (num_rels, num_bases). It should be set to ``None``
+            when ``regularizer=None``.
path: python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py
text: @user1:
I am a little bit confused here. What will happen in regularizer is not "None" and coeff is also not "None"?

The CuGraphRelGraphConv calls RelGraphConvAgg, and regularizer is a property in CuGraphRelGraphConv. From the code structure POV, it is a little bit weird that a variable in a callee has a restriction depends on caller.


@author:
I have updated the doc to avoid ambiguity here. `RelGraphConvAgg` does not use the `RelGraphConvAgg` variable at all, as the regularization information is carried by coeff itself.

diff_hunk: @@ -0,0 +1,388 @@
+"""Torch Module for Relational graph convolution layer using the
+aggregation primitives in cugraph-ops"""
+# pylint: disable=no-member, arguments-differ, invalid-name, too-many-arguments
+import math
+
+import torch as th
+from torch import nn
+
+try:
+    from pylibcugraphops.legacy.aggregators.node_level import (
+        agg_hg_basis_post_bwd_int32,
+        agg_hg_basis_post_bwd_int64,
+        agg_hg_basis_post_fwd_int32,
+        agg_hg_basis_post_fwd_int64,
+    )
+    from pylibcugraphops.legacy.structure.graph_types import (
+        message_flow_graph_hg_csr_int32,
+        message_flow_graph_hg_csr_int64,
+    )
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "dgl.nn.CuGraphRelGraphConv requires pylibcugraphops to be installed."
+    )
+
+CUDA_SM_PER_BLOCK = 49152
+
+
+class RelGraphConvAgg(th.autograd.Function):
+    r"""Custom autograd function for R-GCN aggregation layer that uses the
+    aggregation functions in cugraph-ops."""
+
+    @staticmethod
+    def forward(ctx, g, fanout, num_rels, edge_types, feat, coeff):
+        r"""Compute the forward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        fanout : int
+            Maximum in-degree of nodes.
+        num_rels : int
+            Number of relations.
+        edge_types : torch.Tensor
+            A 1D tensor of edge types. Note that pylibcugraphops only accepts
+            edge_types in int32, so any input of other integer types will be casted
+            into int32, thus introducing some overhead. Pass in int32 tensor for
+            best performance.
+        feat : torch.Tensor, dtype=torch.float32, requires_grad=True
+            Tensor of source node features. Shape: (num_src_nodes, in_feat).
+        coeff : torch.Tensor, dtype=torch.float32, requires_grad=True
+            The coefficient matrix used in basis-decomposition regularization.
+            Shape: (num_rels, num_bases). It should be set to ``None``
+            when ``regularizer=None``.
+
+        Returns
+        -------
+        agg_output : torch.Tensor, dtype=torch.float32
+            Aggregation output. Shape: (num_dst_nodes, num_rels * in_feat)
+            when ``regularizer=None``; Shape: (num_dst_nodes, num_bases * in_feat)
+            when ``regularizer='basis'``.
+        """
+        ctx.device = g.device
+        if g.idtype == th.int32:
+            mfg_csr_func = message_flow_graph_hg_csr_int32
+            agg_fwd_func = agg_hg_basis_post_fwd_int32
+        elif g.idtype == th.int64:
+            mfg_csr_func = message_flow_graph_hg_csr_int64
+            agg_fwd_func = agg_hg_basis_post_fwd_int64
+        else:
+            raise TypeError(
+                f"Supported ID type: torch.int32 or torch.int64, "
+                f"but got {g.idtype}"
+            )
+        ctx.idtype = g.idtype
+
+        _in_feat = feat.shape[-1]
+        indptr, indices, edge_ids = g.adj_sparse("csc")
+        # edge_ids is in a mixed order, need to permutate incoming etypes
+        # and stash the result for backward propagation
+        ctx.edge_types_int32 = edge_types[edge_ids.long()].int()
+        # node_types are not needed by the post-variant RGCN aggregation functions
+        _num_node_types = 0
+        _out_node_types = _in_node_types = None
+
+        mfg = mfg_csr_func(
+            fanout,
+            g.dstnodes(),
+            g.srcnodes(),
+            indptr,
+            indices,
+            _num_node_types,
+            num_rels,
+            _out_node_types,
+            _in_node_types,
+            ctx.edge_types_int32,
+        )
+        ctx.mfg = mfg
+
+        if coeff is None:
+            leading_dimension = num_rels * _in_feat
+        else:
+            _num_bases = coeff.shape[-1]
+            leading_dimension = _num_bases * _in_feat
+
+        agg_output = th.empty(
+            g.num_dst_nodes(),
+            leading_dimension,
+            dtype=th.float32,
+            device=ctx.device,
+        )
+        if coeff is None:
+            agg_fwd_func(agg_output, feat.detach(), mfg)
+        else:
+            agg_fwd_func(
+                agg_output,
+                feat.detach(),
+                mfg,
+                weights_combination=coeff.detach()
+            )
+
+        ctx.save_for_backward(feat, coeff)
+        return agg_output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        r"""Compute the backward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        grad_output : torch.Tensor, dtype=torch.float32
+            Gradient of loss function w.r.t output.
+        """
+        feat, coeff = ctx.saved_tensors
+
+        if ctx.idtype == th.int32:
+            agg_bwd_func = agg_hg_basis_post_bwd_int32
+        else:
+            agg_bwd_func = agg_hg_basis_post_bwd_int64
+
+        grad_feat = th.empty_like(feat, dtype=th.float32, device=ctx.device)
+        if coeff is None:
+            grad_coeff = None
+            agg_bwd_func(grad_feat, grad_output, feat.detach(), ctx.mfg)
+        else:
+            grad_coeff = th.empty_like(
+                coeff, dtype=th.float32, device=ctx.device
+            )
+            agg_bwd_func(
+                grad_feat,
+                grad_output,
+                feat.detach(),
+                ctx.mfg,
+                output_weight_gradient=grad_coeff,
+                weights_combination=coeff.detach()
+            )
+
+        return None, None, None, None, grad_feat, grad_coeff
+
+
+class CuGraphRelGraphConv(nn.Module):
+    r"""An accelerated relational graph convolution layer using the
+    highly-optimized aggregation primitives in cugraph-ops.
+    See :class:`dgl.nn.pytorch.conv.RelGraphConv` for mathematical model.
+
+    .. note::
+        This is an **experimental** feature.
+        Compared with :class:`dgl.nn.pytorch.conv.RelGraphConv`, this model:
+
+        * Only works on cuda devices.
+        * Only support basis-decomposition regularization.
path: python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py
text: @user1:
nit: support -> supports

diff_hunk: @@ -0,0 +1,388 @@
+"""Torch Module for Relational graph convolution layer using the
+aggregation primitives in cugraph-ops"""
+# pylint: disable=no-member, arguments-differ, invalid-name, too-many-arguments
+import math
+
+import torch as th
+from torch import nn
+
+try:
+    from pylibcugraphops.legacy.aggregators.node_level import (
+        agg_hg_basis_post_bwd_int32,
+        agg_hg_basis_post_bwd_int64,
+        agg_hg_basis_post_fwd_int32,
+        agg_hg_basis_post_fwd_int64,
+    )
+    from pylibcugraphops.legacy.structure.graph_types import (
+        message_flow_graph_hg_csr_int32,
+        message_flow_graph_hg_csr_int64,
+    )
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "dgl.nn.CuGraphRelGraphConv requires pylibcugraphops to be installed."
+    )
+
+CUDA_SM_PER_BLOCK = 49152
+
+
+class RelGraphConvAgg(th.autograd.Function):
+    r"""Custom autograd function for R-GCN aggregation layer that uses the
+    aggregation functions in cugraph-ops."""
+
+    @staticmethod
+    def forward(ctx, g, fanout, num_rels, edge_types, feat, coeff):
+        r"""Compute the forward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        fanout : int
+            Maximum in-degree of nodes.
+        num_rels : int
+            Number of relations.
+        edge_types : torch.Tensor
+            A 1D tensor of edge types. Note that pylibcugraphops only accepts
+            edge_types in int32, so any input of other integer types will be casted
+            into int32, thus introducing some overhead. Pass in int32 tensor for
+            best performance.
+        feat : torch.Tensor, dtype=torch.float32, requires_grad=True
+            Tensor of source node features. Shape: (num_src_nodes, in_feat).
+        coeff : torch.Tensor, dtype=torch.float32, requires_grad=True
+            The coefficient matrix used in basis-decomposition regularization.
+            Shape: (num_rels, num_bases). It should be set to ``None``
+            when ``regularizer=None``.
+
+        Returns
+        -------
+        agg_output : torch.Tensor, dtype=torch.float32
+            Aggregation output. Shape: (num_dst_nodes, num_rels * in_feat)
+            when ``regularizer=None``; Shape: (num_dst_nodes, num_bases * in_feat)
+            when ``regularizer='basis'``.
+        """
+        ctx.device = g.device
+        if g.idtype == th.int32:
+            mfg_csr_func = message_flow_graph_hg_csr_int32
+            agg_fwd_func = agg_hg_basis_post_fwd_int32
+        elif g.idtype == th.int64:
+            mfg_csr_func = message_flow_graph_hg_csr_int64
+            agg_fwd_func = agg_hg_basis_post_fwd_int64
+        else:
+            raise TypeError(
+                f"Supported ID type: torch.int32 or torch.int64, "
+                f"but got {g.idtype}"
+            )
+        ctx.idtype = g.idtype
+
+        _in_feat = feat.shape[-1]
+        indptr, indices, edge_ids = g.adj_sparse("csc")
+        # edge_ids is in a mixed order, need to permutate incoming etypes
+        # and stash the result for backward propagation
+        ctx.edge_types_int32 = edge_types[edge_ids.long()].int()
+        # node_types are not needed by the post-variant RGCN aggregation functions
+        _num_node_types = 0
+        _out_node_types = _in_node_types = None
+
+        mfg = mfg_csr_func(
+            fanout,
+            g.dstnodes(),
+            g.srcnodes(),
+            indptr,
+            indices,
+            _num_node_types,
+            num_rels,
+            _out_node_types,
+            _in_node_types,
+            ctx.edge_types_int32,
+        )
+        ctx.mfg = mfg
+
+        if coeff is None:
+            leading_dimension = num_rels * _in_feat
+        else:
+            _num_bases = coeff.shape[-1]
+            leading_dimension = _num_bases * _in_feat
+
+        agg_output = th.empty(
+            g.num_dst_nodes(),
+            leading_dimension,
+            dtype=th.float32,
+            device=ctx.device,
+        )
+        if coeff is None:
+            agg_fwd_func(agg_output, feat.detach(), mfg)
+        else:
+            agg_fwd_func(
+                agg_output,
+                feat.detach(),
+                mfg,
+                weights_combination=coeff.detach()
+            )
+
+        ctx.save_for_backward(feat, coeff)
+        return agg_output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        r"""Compute the backward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        grad_output : torch.Tensor, dtype=torch.float32
+            Gradient of loss function w.r.t output.
+        """
+        feat, coeff = ctx.saved_tensors
+
+        if ctx.idtype == th.int32:
+            agg_bwd_func = agg_hg_basis_post_bwd_int32
+        else:
+            agg_bwd_func = agg_hg_basis_post_bwd_int64
+
+        grad_feat = th.empty_like(feat, dtype=th.float32, device=ctx.device)
+        if coeff is None:
+            grad_coeff = None
+            agg_bwd_func(grad_feat, grad_output, feat.detach(), ctx.mfg)
+        else:
+            grad_coeff = th.empty_like(
+                coeff, dtype=th.float32, device=ctx.device
+            )
+            agg_bwd_func(
+                grad_feat,
+                grad_output,
+                feat.detach(),
+                ctx.mfg,
+                output_weight_gradient=grad_coeff,
+                weights_combination=coeff.detach()
+            )
+
+        return None, None, None, None, grad_feat, grad_coeff
+
+
+class CuGraphRelGraphConv(nn.Module):
+    r"""An accelerated relational graph convolution layer using the
+    highly-optimized aggregation primitives in cugraph-ops.
+    See :class:`dgl.nn.pytorch.conv.RelGraphConv` for mathematical model.
+
+    .. note::
+        This is an **experimental** feature.
+        Compared with :class:`dgl.nn.pytorch.conv.RelGraphConv`, this model:
+
+        * Only works on cuda devices.
+        * Only support basis-decomposition regularization.
+        * Requires an extra argument `fanout` as input, since the current model is
+          designed for sampled-graph (message-flow-graph) use-cases.
+          Full-graph support will be added in upcoming releases.
+
+    Parameters
+    ----------
+    in_feat : int
+        Input feature size.
+    out_feat : int
+        Output feature size.
+    num_rels : int
+        Number of relations.
+    fanout : int
path: python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py
text: @user1:
fanout -> max_in_degree

Can we make this argument optional? E.g., if users do not specify this, calculate it from the input graph?

@author:
Fixed.

diff_hunk: @@ -0,0 +1,388 @@
+"""Torch Module for Relational graph convolution layer using the
+aggregation primitives in cugraph-ops"""
+# pylint: disable=no-member, arguments-differ, invalid-name, too-many-arguments
+import math
+
+import torch as th
+from torch import nn
+
+try:
+    from pylibcugraphops.legacy.aggregators.node_level import (
+        agg_hg_basis_post_bwd_int32,
+        agg_hg_basis_post_bwd_int64,
+        agg_hg_basis_post_fwd_int32,
+        agg_hg_basis_post_fwd_int64,
+    )
+    from pylibcugraphops.legacy.structure.graph_types import (
+        message_flow_graph_hg_csr_int32,
+        message_flow_graph_hg_csr_int64,
+    )
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "dgl.nn.CuGraphRelGraphConv requires pylibcugraphops to be installed."
+    )
+
+CUDA_SM_PER_BLOCK = 49152
+
+
+class RelGraphConvAgg(th.autograd.Function):
+    r"""Custom autograd function for R-GCN aggregation layer that uses the
+    aggregation functions in cugraph-ops."""
+
+    @staticmethod
+    def forward(ctx, g, fanout, num_rels, edge_types, feat, coeff):
+        r"""Compute the forward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        fanout : int
+            Maximum in-degree of nodes.
+        num_rels : int
+            Number of relations.
+        edge_types : torch.Tensor
+            A 1D tensor of edge types. Note that pylibcugraphops only accepts
+            edge_types in int32, so any input of other integer types will be casted
+            into int32, thus introducing some overhead. Pass in int32 tensor for
+            best performance.
+        feat : torch.Tensor, dtype=torch.float32, requires_grad=True
+            Tensor of source node features. Shape: (num_src_nodes, in_feat).
+        coeff : torch.Tensor, dtype=torch.float32, requires_grad=True
+            The coefficient matrix used in basis-decomposition regularization.
+            Shape: (num_rels, num_bases). It should be set to ``None``
+            when ``regularizer=None``.
+
+        Returns
+        -------
+        agg_output : torch.Tensor, dtype=torch.float32
+            Aggregation output. Shape: (num_dst_nodes, num_rels * in_feat)
+            when ``regularizer=None``; Shape: (num_dst_nodes, num_bases * in_feat)
+            when ``regularizer='basis'``.
+        """
+        ctx.device = g.device
+        if g.idtype == th.int32:
+            mfg_csr_func = message_flow_graph_hg_csr_int32
+            agg_fwd_func = agg_hg_basis_post_fwd_int32
+        elif g.idtype == th.int64:
+            mfg_csr_func = message_flow_graph_hg_csr_int64
+            agg_fwd_func = agg_hg_basis_post_fwd_int64
+        else:
+            raise TypeError(
+                f"Supported ID type: torch.int32 or torch.int64, "
+                f"but got {g.idtype}"
+            )
+        ctx.idtype = g.idtype
+
+        _in_feat = feat.shape[-1]
+        indptr, indices, edge_ids = g.adj_sparse("csc")
+        # edge_ids is in a mixed order, need to permutate incoming etypes
+        # and stash the result for backward propagation
+        ctx.edge_types_int32 = edge_types[edge_ids.long()].int()
+        # node_types are not needed by the post-variant RGCN aggregation functions
+        _num_node_types = 0
+        _out_node_types = _in_node_types = None
+
+        mfg = mfg_csr_func(
+            fanout,
+            g.dstnodes(),
+            g.srcnodes(),
+            indptr,
+            indices,
+            _num_node_types,
+            num_rels,
+            _out_node_types,
+            _in_node_types,
+            ctx.edge_types_int32,
+        )
+        ctx.mfg = mfg
+
+        if coeff is None:
+            leading_dimension = num_rels * _in_feat
+        else:
+            _num_bases = coeff.shape[-1]
+            leading_dimension = _num_bases * _in_feat
+
+        agg_output = th.empty(
+            g.num_dst_nodes(),
+            leading_dimension,
+            dtype=th.float32,
+            device=ctx.device,
+        )
+        if coeff is None:
+            agg_fwd_func(agg_output, feat.detach(), mfg)
+        else:
+            agg_fwd_func(
+                agg_output,
+                feat.detach(),
+                mfg,
+                weights_combination=coeff.detach()
+            )
+
+        ctx.save_for_backward(feat, coeff)
+        return agg_output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        r"""Compute the backward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        grad_output : torch.Tensor, dtype=torch.float32
+            Gradient of loss function w.r.t output.
+        """
+        feat, coeff = ctx.saved_tensors
+
+        if ctx.idtype == th.int32:
+            agg_bwd_func = agg_hg_basis_post_bwd_int32
+        else:
+            agg_bwd_func = agg_hg_basis_post_bwd_int64
+
+        grad_feat = th.empty_like(feat, dtype=th.float32, device=ctx.device)
+        if coeff is None:
+            grad_coeff = None
+            agg_bwd_func(grad_feat, grad_output, feat.detach(), ctx.mfg)
+        else:
+            grad_coeff = th.empty_like(
+                coeff, dtype=th.float32, device=ctx.device
+            )
+            agg_bwd_func(
+                grad_feat,
+                grad_output,
+                feat.detach(),
+                ctx.mfg,
+                output_weight_gradient=grad_coeff,
+                weights_combination=coeff.detach()
+            )
+
+        return None, None, None, None, grad_feat, grad_coeff
+
+
+class CuGraphRelGraphConv(nn.Module):
+    r"""An accelerated relational graph convolution layer using the
+    highly-optimized aggregation primitives in cugraph-ops.
+    See :class:`dgl.nn.pytorch.conv.RelGraphConv` for mathematical model.
+
+    .. note::
+        This is an **experimental** feature.
+        Compared with :class:`dgl.nn.pytorch.conv.RelGraphConv`, this model:
path: python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py
text: @user1:
I think you can remove L168.

@user2:
I think this module is a special case that it is a direct optimization on our `RelGraphConv`. We can expect the users to understand the original `RelGraphConv` first, so making docstring self-contained is less a concern.

diff_hunk: @@ -0,0 +1,388 @@
+"""Torch Module for Relational graph convolution layer using the
+aggregation primitives in cugraph-ops"""
+# pylint: disable=no-member, arguments-differ, invalid-name, too-many-arguments
+import math
+
+import torch as th
+from torch import nn
+
+try:
+    from pylibcugraphops.legacy.aggregators.node_level import (
+        agg_hg_basis_post_bwd_int32,
+        agg_hg_basis_post_bwd_int64,
+        agg_hg_basis_post_fwd_int32,
+        agg_hg_basis_post_fwd_int64,
+    )
+    from pylibcugraphops.legacy.structure.graph_types import (
+        message_flow_graph_hg_csr_int32,
+        message_flow_graph_hg_csr_int64,
+    )
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "dgl.nn.CuGraphRelGraphConv requires pylibcugraphops to be installed."
+    )
+
+CUDA_SM_PER_BLOCK = 49152
+
+
+class RelGraphConvAgg(th.autograd.Function):
+    r"""Custom autograd function for R-GCN aggregation layer that uses the
+    aggregation functions in cugraph-ops."""
+
+    @staticmethod
+    def forward(ctx, g, fanout, num_rels, edge_types, feat, coeff):
+        r"""Compute the forward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        fanout : int
+            Maximum in-degree of nodes.
+        num_rels : int
+            Number of relations.
+        edge_types : torch.Tensor
+            A 1D tensor of edge types. Note that pylibcugraphops only accepts
+            edge_types in int32, so any input of other integer types will be casted
+            into int32, thus introducing some overhead. Pass in int32 tensor for
+            best performance.
+        feat : torch.Tensor, dtype=torch.float32, requires_grad=True
+            Tensor of source node features. Shape: (num_src_nodes, in_feat).
+        coeff : torch.Tensor, dtype=torch.float32, requires_grad=True
+            The coefficient matrix used in basis-decomposition regularization.
+            Shape: (num_rels, num_bases). It should be set to ``None``
+            when ``regularizer=None``.
+
+        Returns
+        -------
+        agg_output : torch.Tensor, dtype=torch.float32
+            Aggregation output. Shape: (num_dst_nodes, num_rels * in_feat)
+            when ``regularizer=None``; Shape: (num_dst_nodes, num_bases * in_feat)
+            when ``regularizer='basis'``.
+        """
+        ctx.device = g.device
+        if g.idtype == th.int32:
+            mfg_csr_func = message_flow_graph_hg_csr_int32
+            agg_fwd_func = agg_hg_basis_post_fwd_int32
+        elif g.idtype == th.int64:
+            mfg_csr_func = message_flow_graph_hg_csr_int64
+            agg_fwd_func = agg_hg_basis_post_fwd_int64
+        else:
+            raise TypeError(
+                f"Supported ID type: torch.int32 or torch.int64, "
+                f"but got {g.idtype}"
+            )
+        ctx.idtype = g.idtype
+
+        _in_feat = feat.shape[-1]
+        indptr, indices, edge_ids = g.adj_sparse("csc")
+        # edge_ids is in a mixed order, need to permutate incoming etypes
+        # and stash the result for backward propagation
+        ctx.edge_types_int32 = edge_types[edge_ids.long()].int()
+        # node_types are not needed by the post-variant RGCN aggregation functions
+        _num_node_types = 0
+        _out_node_types = _in_node_types = None
+
+        mfg = mfg_csr_func(
+            fanout,
+            g.dstnodes(),
+            g.srcnodes(),
+            indptr,
+            indices,
+            _num_node_types,
+            num_rels,
+            _out_node_types,
+            _in_node_types,
+            ctx.edge_types_int32,
+        )
+        ctx.mfg = mfg
+
+        if coeff is None:
+            leading_dimension = num_rels * _in_feat
+        else:
+            _num_bases = coeff.shape[-1]
+            leading_dimension = _num_bases * _in_feat
+
+        agg_output = th.empty(
+            g.num_dst_nodes(),
+            leading_dimension,
+            dtype=th.float32,
+            device=ctx.device,
+        )
+        if coeff is None:
+            agg_fwd_func(agg_output, feat.detach(), mfg)
+        else:
+            agg_fwd_func(
+                agg_output,
+                feat.detach(),
+                mfg,
+                weights_combination=coeff.detach()
+            )
+
+        ctx.save_for_backward(feat, coeff)
+        return agg_output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        r"""Compute the backward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        grad_output : torch.Tensor, dtype=torch.float32
+            Gradient of loss function w.r.t output.
+        """
+        feat, coeff = ctx.saved_tensors
+
+        if ctx.idtype == th.int32:
+            agg_bwd_func = agg_hg_basis_post_bwd_int32
+        else:
+            agg_bwd_func = agg_hg_basis_post_bwd_int64
+
+        grad_feat = th.empty_like(feat, dtype=th.float32, device=ctx.device)
+        if coeff is None:
+            grad_coeff = None
+            agg_bwd_func(grad_feat, grad_output, feat.detach(), ctx.mfg)
+        else:
+            grad_coeff = th.empty_like(
+                coeff, dtype=th.float32, device=ctx.device
+            )
+            agg_bwd_func(
+                grad_feat,
+                grad_output,
+                feat.detach(),
+                ctx.mfg,
+                output_weight_gradient=grad_coeff,
+                weights_combination=coeff.detach()
+            )
+
+        return None, None, None, None, grad_feat, grad_coeff
+
+
+class CuGraphRelGraphConv(nn.Module):
+    r"""An accelerated relational graph convolution layer using the
+    highly-optimized aggregation primitives in cugraph-ops.
+    See :class:`dgl.nn.pytorch.conv.RelGraphConv` for mathematical model.
+
+    .. note::
+        This is an **experimental** feature.
+        Compared with :class:`dgl.nn.pytorch.conv.RelGraphConv`, this model:
+
+        * Only works on cuda devices.
+        * Only support basis-decomposition regularization.
+        * Requires an extra argument `fanout` as input, since the current model is
+          designed for sampled-graph (message-flow-graph) use-cases.
+          Full-graph support will be added in upcoming releases.
path: python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py
text: @user1:
I suggest replacing L172-174 with

> Only supports message flow graphs and not full graphs

- You have a dedicated description of the `fanout` argument below, so you do not need to mention that here.
- "Full-graph support will be added in upcoming releases." conveys little information and it's better to not mention that at all. If users want that, they will open a feature request.

@author:
Agree. This is redundant.

diff_hunk: @@ -0,0 +1,388 @@
+"""Torch Module for Relational graph convolution layer using the
+aggregation primitives in cugraph-ops"""
+# pylint: disable=no-member, arguments-differ, invalid-name, too-many-arguments
+import math
+
+import torch as th
+from torch import nn
+
+try:
+    from pylibcugraphops.legacy.aggregators.node_level import (
+        agg_hg_basis_post_bwd_int32,
+        agg_hg_basis_post_bwd_int64,
+        agg_hg_basis_post_fwd_int32,
+        agg_hg_basis_post_fwd_int64,
+    )
+    from pylibcugraphops.legacy.structure.graph_types import (
+        message_flow_graph_hg_csr_int32,
+        message_flow_graph_hg_csr_int64,
+    )
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "dgl.nn.CuGraphRelGraphConv requires pylibcugraphops to be installed."
+    )
+
+CUDA_SM_PER_BLOCK = 49152
+
+
+class RelGraphConvAgg(th.autograd.Function):
+    r"""Custom autograd function for R-GCN aggregation layer that uses the
+    aggregation functions in cugraph-ops."""
+
+    @staticmethod
+    def forward(ctx, g, fanout, num_rels, edge_types, feat, coeff):
+        r"""Compute the forward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        fanout : int
+            Maximum in-degree of nodes.
+        num_rels : int
+            Number of relations.
+        edge_types : torch.Tensor
+            A 1D tensor of edge types. Note that pylibcugraphops only accepts
+            edge_types in int32, so any input of other integer types will be casted
+            into int32, thus introducing some overhead. Pass in int32 tensor for
+            best performance.
+        feat : torch.Tensor, dtype=torch.float32, requires_grad=True
+            Tensor of source node features. Shape: (num_src_nodes, in_feat).
+        coeff : torch.Tensor, dtype=torch.float32, requires_grad=True
+            The coefficient matrix used in basis-decomposition regularization.
+            Shape: (num_rels, num_bases). It should be set to ``None``
+            when ``regularizer=None``.
+
+        Returns
+        -------
+        agg_output : torch.Tensor, dtype=torch.float32
+            Aggregation output. Shape: (num_dst_nodes, num_rels * in_feat)
+            when ``regularizer=None``; Shape: (num_dst_nodes, num_bases * in_feat)
+            when ``regularizer='basis'``.
+        """
+        ctx.device = g.device
+        if g.idtype == th.int32:
+            mfg_csr_func = message_flow_graph_hg_csr_int32
+            agg_fwd_func = agg_hg_basis_post_fwd_int32
+        elif g.idtype == th.int64:
+            mfg_csr_func = message_flow_graph_hg_csr_int64
+            agg_fwd_func = agg_hg_basis_post_fwd_int64
+        else:
+            raise TypeError(
+                f"Supported ID type: torch.int32 or torch.int64, "
+                f"but got {g.idtype}"
+            )
+        ctx.idtype = g.idtype
+
+        _in_feat = feat.shape[-1]
+        indptr, indices, edge_ids = g.adj_sparse("csc")
+        # edge_ids is in a mixed order, need to permutate incoming etypes
+        # and stash the result for backward propagation
+        ctx.edge_types_int32 = edge_types[edge_ids.long()].int()
+        # node_types are not needed by the post-variant RGCN aggregation functions
+        _num_node_types = 0
+        _out_node_types = _in_node_types = None
+
+        mfg = mfg_csr_func(
+            fanout,
+            g.dstnodes(),
+            g.srcnodes(),
+            indptr,
+            indices,
+            _num_node_types,
+            num_rels,
+            _out_node_types,
+            _in_node_types,
+            ctx.edge_types_int32,
+        )
+        ctx.mfg = mfg
+
+        if coeff is None:
+            leading_dimension = num_rels * _in_feat
+        else:
+            _num_bases = coeff.shape[-1]
+            leading_dimension = _num_bases * _in_feat
+
+        agg_output = th.empty(
+            g.num_dst_nodes(),
+            leading_dimension,
+            dtype=th.float32,
+            device=ctx.device,
+        )
+        if coeff is None:
+            agg_fwd_func(agg_output, feat.detach(), mfg)
+        else:
+            agg_fwd_func(
+                agg_output,
+                feat.detach(),
+                mfg,
+                weights_combination=coeff.detach()
+            )
+
+        ctx.save_for_backward(feat, coeff)
+        return agg_output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        r"""Compute the backward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        grad_output : torch.Tensor, dtype=torch.float32
+            Gradient of loss function w.r.t output.
+        """
+        feat, coeff = ctx.saved_tensors
+
+        if ctx.idtype == th.int32:
+            agg_bwd_func = agg_hg_basis_post_bwd_int32
+        else:
+            agg_bwd_func = agg_hg_basis_post_bwd_int64
+
+        grad_feat = th.empty_like(feat, dtype=th.float32, device=ctx.device)
+        if coeff is None:
+            grad_coeff = None
+            agg_bwd_func(grad_feat, grad_output, feat.detach(), ctx.mfg)
+        else:
+            grad_coeff = th.empty_like(
+                coeff, dtype=th.float32, device=ctx.device
+            )
+            agg_bwd_func(
+                grad_feat,
+                grad_output,
+                feat.detach(),
+                ctx.mfg,
+                output_weight_gradient=grad_coeff,
+                weights_combination=coeff.detach()
+            )
+
+        return None, None, None, None, grad_feat, grad_coeff
+
+
+class CuGraphRelGraphConv(nn.Module):
+    r"""An accelerated relational graph convolution layer using the
+    highly-optimized aggregation primitives in cugraph-ops.
+    See :class:`dgl.nn.pytorch.conv.RelGraphConv` for mathematical model.
+
+    .. note::
+        This is an **experimental** feature.
+        Compared with :class:`dgl.nn.pytorch.conv.RelGraphConv`, this model:
+
+        * Only works on cuda devices.
+        * Only support basis-decomposition regularization.
+        * Requires an extra argument `fanout` as input, since the current model is
+          designed for sampled-graph (message-flow-graph) use-cases.
+          Full-graph support will be added in upcoming releases.
+
+    Parameters
+    ----------
+    in_feat : int
+        Input feature size.
+    out_feat : int
+        Output feature size.
+    num_rels : int
+        Number of relations.
+    fanout : int
+        Maximum number of sampled neighbors of an destination node;
path: python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py
text: @user1:
```suggestion
        Maximum number of sampled neighbors of a destination node,
```

diff_hunk: @@ -0,0 +1,388 @@
+"""Torch Module for Relational graph convolution layer using the
+aggregation primitives in cugraph-ops"""
+# pylint: disable=no-member, arguments-differ, invalid-name, too-many-arguments
+import math
+
+import torch as th
+from torch import nn
+
+try:
+    from pylibcugraphops.legacy.aggregators.node_level import (
+        agg_hg_basis_post_bwd_int32,
+        agg_hg_basis_post_bwd_int64,
+        agg_hg_basis_post_fwd_int32,
+        agg_hg_basis_post_fwd_int64,
+    )
+    from pylibcugraphops.legacy.structure.graph_types import (
+        message_flow_graph_hg_csr_int32,
+        message_flow_graph_hg_csr_int64,
+    )
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "dgl.nn.CuGraphRelGraphConv requires pylibcugraphops to be installed."
+    )
+
+CUDA_SM_PER_BLOCK = 49152
+
+
+class RelGraphConvAgg(th.autograd.Function):
+    r"""Custom autograd function for R-GCN aggregation layer that uses the
+    aggregation functions in cugraph-ops."""
+
+    @staticmethod
+    def forward(ctx, g, fanout, num_rels, edge_types, feat, coeff):
+        r"""Compute the forward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        fanout : int
+            Maximum in-degree of nodes.
+        num_rels : int
+            Number of relations.
+        edge_types : torch.Tensor
+            A 1D tensor of edge types. Note that pylibcugraphops only accepts
+            edge_types in int32, so any input of other integer types will be casted
+            into int32, thus introducing some overhead. Pass in int32 tensor for
+            best performance.
+        feat : torch.Tensor, dtype=torch.float32, requires_grad=True
+            Tensor of source node features. Shape: (num_src_nodes, in_feat).
+        coeff : torch.Tensor, dtype=torch.float32, requires_grad=True
+            The coefficient matrix used in basis-decomposition regularization.
+            Shape: (num_rels, num_bases). It should be set to ``None``
+            when ``regularizer=None``.
+
+        Returns
+        -------
+        agg_output : torch.Tensor, dtype=torch.float32
+            Aggregation output. Shape: (num_dst_nodes, num_rels * in_feat)
+            when ``regularizer=None``; Shape: (num_dst_nodes, num_bases * in_feat)
+            when ``regularizer='basis'``.
+        """
+        ctx.device = g.device
+        if g.idtype == th.int32:
+            mfg_csr_func = message_flow_graph_hg_csr_int32
+            agg_fwd_func = agg_hg_basis_post_fwd_int32
+        elif g.idtype == th.int64:
+            mfg_csr_func = message_flow_graph_hg_csr_int64
+            agg_fwd_func = agg_hg_basis_post_fwd_int64
+        else:
+            raise TypeError(
+                f"Supported ID type: torch.int32 or torch.int64, "
+                f"but got {g.idtype}"
+            )
+        ctx.idtype = g.idtype
+
+        _in_feat = feat.shape[-1]
+        indptr, indices, edge_ids = g.adj_sparse("csc")
+        # edge_ids is in a mixed order, need to permutate incoming etypes
+        # and stash the result for backward propagation
+        ctx.edge_types_int32 = edge_types[edge_ids.long()].int()
+        # node_types are not needed by the post-variant RGCN aggregation functions
+        _num_node_types = 0
+        _out_node_types = _in_node_types = None
+
+        mfg = mfg_csr_func(
+            fanout,
+            g.dstnodes(),
+            g.srcnodes(),
+            indptr,
+            indices,
+            _num_node_types,
+            num_rels,
+            _out_node_types,
+            _in_node_types,
+            ctx.edge_types_int32,
+        )
+        ctx.mfg = mfg
+
+        if coeff is None:
+            leading_dimension = num_rels * _in_feat
+        else:
+            _num_bases = coeff.shape[-1]
+            leading_dimension = _num_bases * _in_feat
+
+        agg_output = th.empty(
+            g.num_dst_nodes(),
+            leading_dimension,
+            dtype=th.float32,
+            device=ctx.device,
+        )
+        if coeff is None:
+            agg_fwd_func(agg_output, feat.detach(), mfg)
+        else:
+            agg_fwd_func(
+                agg_output,
+                feat.detach(),
+                mfg,
+                weights_combination=coeff.detach()
+            )
+
+        ctx.save_for_backward(feat, coeff)
+        return agg_output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        r"""Compute the backward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        grad_output : torch.Tensor, dtype=torch.float32
+            Gradient of loss function w.r.t output.
+        """
+        feat, coeff = ctx.saved_tensors
+
+        if ctx.idtype == th.int32:
+            agg_bwd_func = agg_hg_basis_post_bwd_int32
+        else:
+            agg_bwd_func = agg_hg_basis_post_bwd_int64
+
+        grad_feat = th.empty_like(feat, dtype=th.float32, device=ctx.device)
+        if coeff is None:
+            grad_coeff = None
+            agg_bwd_func(grad_feat, grad_output, feat.detach(), ctx.mfg)
+        else:
+            grad_coeff = th.empty_like(
+                coeff, dtype=th.float32, device=ctx.device
+            )
+            agg_bwd_func(
+                grad_feat,
+                grad_output,
+                feat.detach(),
+                ctx.mfg,
+                output_weight_gradient=grad_coeff,
+                weights_combination=coeff.detach()
+            )
+
+        return None, None, None, None, grad_feat, grad_coeff
+
+
+class CuGraphRelGraphConv(nn.Module):
path: python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py
text: @user2:
Can we change from `pytorch/conv/cugraph_relgraphconv.py/CuGraphRelGraphConv` to `pytorch/conv/cugraph/relgraphconv.py/RelGraphConv`? I just find this name a bit tedious. Also we might want to handle all dependency checks for cugraph once in the same subdirectory. @user1 

@user1:
What's the confusion here?

@user2:
First, we have dependency checks like 

```
try:
    from pylibcugraphops.legacy.aggregators.node_level import (
        agg_hg_basis_post_bwd_int32,
        agg_hg_basis_post_bwd_int64,
        agg_hg_basis_post_fwd_int32,
        agg_hg_basis_post_fwd_int64,
    )
    from pylibcugraphops.legacy.structure.graph_types import (
        message_flow_graph_hg_csr_int32,
        message_flow_graph_hg_csr_int64,
    )
except ModuleNotFoundError:
    raise ModuleNotFoundError(
        "dgl.nn.CuGraphRelGraphConv requires pylibcugraphops to be installed."
    )
```

If we have a dedicated subdirectory for CuGraph, then it can be done there rather than in individual files.

Second, I think it's better to do

```
from dgl.nn.cugraph import AConv, BConv, ...
```

than

```
from dgl.nn import CuGraphAConv, CuGraphBConv
```

These are minor issues. I'm fine to wait for a few more modules to see what will happen.

diff_hunk: @@ -0,0 +1,388 @@
+"""Torch Module for Relational graph convolution layer using the
+aggregation primitives in cugraph-ops"""
+# pylint: disable=no-member, arguments-differ, invalid-name, too-many-arguments
+import math
+
+import torch as th
+from torch import nn
+
+try:
+    from pylibcugraphops.legacy.aggregators.node_level import (
+        agg_hg_basis_post_bwd_int32,
+        agg_hg_basis_post_bwd_int64,
+        agg_hg_basis_post_fwd_int32,
+        agg_hg_basis_post_fwd_int64,
+    )
+    from pylibcugraphops.legacy.structure.graph_types import (
+        message_flow_graph_hg_csr_int32,
+        message_flow_graph_hg_csr_int64,
+    )
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "dgl.nn.CuGraphRelGraphConv requires pylibcugraphops to be installed."
+    )
+
+CUDA_SM_PER_BLOCK = 49152
+
+
+class RelGraphConvAgg(th.autograd.Function):
+    r"""Custom autograd function for R-GCN aggregation layer that uses the
+    aggregation functions in cugraph-ops."""
+
+    @staticmethod
+    def forward(ctx, g, fanout, num_rels, edge_types, feat, coeff):
+        r"""Compute the forward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        fanout : int
+            Maximum in-degree of nodes.
+        num_rels : int
+            Number of relations.
+        edge_types : torch.Tensor
+            A 1D tensor of edge types. Note that pylibcugraphops only accepts
+            edge_types in int32, so any input of other integer types will be casted
+            into int32, thus introducing some overhead. Pass in int32 tensor for
+            best performance.
+        feat : torch.Tensor, dtype=torch.float32, requires_grad=True
+            Tensor of source node features. Shape: (num_src_nodes, in_feat).
+        coeff : torch.Tensor, dtype=torch.float32, requires_grad=True
+            The coefficient matrix used in basis-decomposition regularization.
+            Shape: (num_rels, num_bases). It should be set to ``None``
+            when ``regularizer=None``.
+
+        Returns
+        -------
+        agg_output : torch.Tensor, dtype=torch.float32
+            Aggregation output. Shape: (num_dst_nodes, num_rels * in_feat)
+            when ``regularizer=None``; Shape: (num_dst_nodes, num_bases * in_feat)
+            when ``regularizer='basis'``.
+        """
+        ctx.device = g.device
+        if g.idtype == th.int32:
+            mfg_csr_func = message_flow_graph_hg_csr_int32
+            agg_fwd_func = agg_hg_basis_post_fwd_int32
+        elif g.idtype == th.int64:
+            mfg_csr_func = message_flow_graph_hg_csr_int64
+            agg_fwd_func = agg_hg_basis_post_fwd_int64
+        else:
+            raise TypeError(
+                f"Supported ID type: torch.int32 or torch.int64, "
+                f"but got {g.idtype}"
+            )
+        ctx.idtype = g.idtype
+
+        _in_feat = feat.shape[-1]
+        indptr, indices, edge_ids = g.adj_sparse("csc")
+        # edge_ids is in a mixed order, need to permutate incoming etypes
+        # and stash the result for backward propagation
+        ctx.edge_types_int32 = edge_types[edge_ids.long()].int()
+        # node_types are not needed by the post-variant RGCN aggregation functions
+        _num_node_types = 0
+        _out_node_types = _in_node_types = None
+
+        mfg = mfg_csr_func(
+            fanout,
+            g.dstnodes(),
+            g.srcnodes(),
+            indptr,
+            indices,
+            _num_node_types,
+            num_rels,
+            _out_node_types,
+            _in_node_types,
+            ctx.edge_types_int32,
+        )
+        ctx.mfg = mfg
+
+        if coeff is None:
+            leading_dimension = num_rels * _in_feat
+        else:
+            _num_bases = coeff.shape[-1]
+            leading_dimension = _num_bases * _in_feat
+
+        agg_output = th.empty(
+            g.num_dst_nodes(),
+            leading_dimension,
+            dtype=th.float32,
+            device=ctx.device,
+        )
+        if coeff is None:
+            agg_fwd_func(agg_output, feat.detach(), mfg)
+        else:
+            agg_fwd_func(
+                agg_output,
+                feat.detach(),
+                mfg,
+                weights_combination=coeff.detach()
+            )
+
+        ctx.save_for_backward(feat, coeff)
+        return agg_output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        r"""Compute the backward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        grad_output : torch.Tensor, dtype=torch.float32
+            Gradient of loss function w.r.t output.
+        """
+        feat, coeff = ctx.saved_tensors
+
+        if ctx.idtype == th.int32:
+            agg_bwd_func = agg_hg_basis_post_bwd_int32
+        else:
+            agg_bwd_func = agg_hg_basis_post_bwd_int64
+
+        grad_feat = th.empty_like(feat, dtype=th.float32, device=ctx.device)
+        if coeff is None:
+            grad_coeff = None
+            agg_bwd_func(grad_feat, grad_output, feat.detach(), ctx.mfg)
+        else:
+            grad_coeff = th.empty_like(
+                coeff, dtype=th.float32, device=ctx.device
+            )
+            agg_bwd_func(
+                grad_feat,
+                grad_output,
+                feat.detach(),
+                ctx.mfg,
+                output_weight_gradient=grad_coeff,
+                weights_combination=coeff.detach()
+            )
+
+        return None, None, None, None, grad_feat, grad_coeff
+
+
+class CuGraphRelGraphConv(nn.Module):
+    r"""An accelerated relational graph convolution layer using the
+    highly-optimized aggregation primitives in cugraph-ops.
+    See :class:`dgl.nn.pytorch.conv.RelGraphConv` for mathematical model.
+
+    .. note::
+        This is an **experimental** feature.
+        Compared with :class:`dgl.nn.pytorch.conv.RelGraphConv`, this model:
+
+        * Only works on cuda devices.
+        * Only support basis-decomposition regularization.
+        * Requires an extra argument `fanout` as input, since the current model is
+          designed for sampled-graph (message-flow-graph) use-cases.
+          Full-graph support will be added in upcoming releases.
+
+    Parameters
+    ----------
+    in_feat : int
+        Input feature size.
+    out_feat : int
+        Output feature size.
+    num_rels : int
+        Number of relations.
+    fanout : int
path: python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py
text: @user1:
Maybe move this argument to the beginning or the end for consistency of argument order with `dgl.nn.pytorch.conv.RelGraphConv`.

@author:
Done, and it becomes optional now.

diff_hunk: @@ -0,0 +1,388 @@
+"""Torch Module for Relational graph convolution layer using the
+aggregation primitives in cugraph-ops"""
+# pylint: disable=no-member, arguments-differ, invalid-name, too-many-arguments
+import math
+
+import torch as th
+from torch import nn
+
+try:
+    from pylibcugraphops.legacy.aggregators.node_level import (
+        agg_hg_basis_post_bwd_int32,
+        agg_hg_basis_post_bwd_int64,
+        agg_hg_basis_post_fwd_int32,
+        agg_hg_basis_post_fwd_int64,
+    )
+    from pylibcugraphops.legacy.structure.graph_types import (
+        message_flow_graph_hg_csr_int32,
+        message_flow_graph_hg_csr_int64,
+    )
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "dgl.nn.CuGraphRelGraphConv requires pylibcugraphops to be installed."
+    )
+
+CUDA_SM_PER_BLOCK = 49152
+
+
+class RelGraphConvAgg(th.autograd.Function):
+    r"""Custom autograd function for R-GCN aggregation layer that uses the
+    aggregation functions in cugraph-ops."""
+
+    @staticmethod
+    def forward(ctx, g, fanout, num_rels, edge_types, feat, coeff):
+        r"""Compute the forward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        fanout : int
+            Maximum in-degree of nodes.
+        num_rels : int
+            Number of relations.
+        edge_types : torch.Tensor
+            A 1D tensor of edge types. Note that pylibcugraphops only accepts
+            edge_types in int32, so any input of other integer types will be casted
+            into int32, thus introducing some overhead. Pass in int32 tensor for
+            best performance.
+        feat : torch.Tensor, dtype=torch.float32, requires_grad=True
+            Tensor of source node features. Shape: (num_src_nodes, in_feat).
+        coeff : torch.Tensor, dtype=torch.float32, requires_grad=True
+            The coefficient matrix used in basis-decomposition regularization.
+            Shape: (num_rels, num_bases). It should be set to ``None``
+            when ``regularizer=None``.
+
+        Returns
+        -------
+        agg_output : torch.Tensor, dtype=torch.float32
+            Aggregation output. Shape: (num_dst_nodes, num_rels * in_feat)
+            when ``regularizer=None``; Shape: (num_dst_nodes, num_bases * in_feat)
+            when ``regularizer='basis'``.
+        """
+        ctx.device = g.device
+        if g.idtype == th.int32:
+            mfg_csr_func = message_flow_graph_hg_csr_int32
+            agg_fwd_func = agg_hg_basis_post_fwd_int32
+        elif g.idtype == th.int64:
+            mfg_csr_func = message_flow_graph_hg_csr_int64
+            agg_fwd_func = agg_hg_basis_post_fwd_int64
+        else:
+            raise TypeError(
+                f"Supported ID type: torch.int32 or torch.int64, "
+                f"but got {g.idtype}"
+            )
+        ctx.idtype = g.idtype
+
+        _in_feat = feat.shape[-1]
+        indptr, indices, edge_ids = g.adj_sparse("csc")
+        # edge_ids is in a mixed order, need to permutate incoming etypes
+        # and stash the result for backward propagation
+        ctx.edge_types_int32 = edge_types[edge_ids.long()].int()
+        # node_types are not needed by the post-variant RGCN aggregation functions
+        _num_node_types = 0
+        _out_node_types = _in_node_types = None
+
+        mfg = mfg_csr_func(
+            fanout,
+            g.dstnodes(),
+            g.srcnodes(),
+            indptr,
+            indices,
+            _num_node_types,
+            num_rels,
+            _out_node_types,
+            _in_node_types,
+            ctx.edge_types_int32,
+        )
+        ctx.mfg = mfg
+
+        if coeff is None:
+            leading_dimension = num_rels * _in_feat
+        else:
+            _num_bases = coeff.shape[-1]
+            leading_dimension = _num_bases * _in_feat
+
+        agg_output = th.empty(
+            g.num_dst_nodes(),
+            leading_dimension,
+            dtype=th.float32,
+            device=ctx.device,
+        )
+        if coeff is None:
+            agg_fwd_func(agg_output, feat.detach(), mfg)
+        else:
+            agg_fwd_func(
+                agg_output,
+                feat.detach(),
+                mfg,
+                weights_combination=coeff.detach()
+            )
+
+        ctx.save_for_backward(feat, coeff)
+        return agg_output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        r"""Compute the backward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        grad_output : torch.Tensor, dtype=torch.float32
+            Gradient of loss function w.r.t output.
+        """
+        feat, coeff = ctx.saved_tensors
+
+        if ctx.idtype == th.int32:
+            agg_bwd_func = agg_hg_basis_post_bwd_int32
+        else:
+            agg_bwd_func = agg_hg_basis_post_bwd_int64
+
+        grad_feat = th.empty_like(feat, dtype=th.float32, device=ctx.device)
+        if coeff is None:
+            grad_coeff = None
+            agg_bwd_func(grad_feat, grad_output, feat.detach(), ctx.mfg)
+        else:
+            grad_coeff = th.empty_like(
+                coeff, dtype=th.float32, device=ctx.device
+            )
+            agg_bwd_func(
+                grad_feat,
+                grad_output,
+                feat.detach(),
+                ctx.mfg,
+                output_weight_gradient=grad_coeff,
+                weights_combination=coeff.detach()
+            )
+
+        return None, None, None, None, grad_feat, grad_coeff
+
+
+class CuGraphRelGraphConv(nn.Module):
+    r"""An accelerated relational graph convolution layer using the
+    highly-optimized aggregation primitives in cugraph-ops.
+    See :class:`dgl.nn.pytorch.conv.RelGraphConv` for mathematical model.
+
+    .. note::
+        This is an **experimental** feature.
+        Compared with :class:`dgl.nn.pytorch.conv.RelGraphConv`, this model:
+
+        * Only works on cuda devices.
+        * Only support basis-decomposition regularization.
+        * Requires an extra argument `fanout` as input, since the current model is
+          designed for sampled-graph (message-flow-graph) use-cases.
+          Full-graph support will be added in upcoming releases.
+
+    Parameters
+    ----------
+    in_feat : int
+        Input feature size.
+    out_feat : int
+        Output feature size.
+    num_rels : int
+        Number of relations.
+    fanout : int
+        Maximum number of sampled neighbors of an destination node;
+        i.e, maximum in degree of destination nodes
path: python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py
text: @user1:
```suggestion
        i.e. maximum in degree of destination nodes
```

This is not necessarily the case and might be dependent on the sampler configuration, right?

@author:
This argument should be renamed as max_in_degree, as Minjie suggested.
If we use `NeighborSampler` to generate sampled graphs, max_in_degree should be equivalent to the fanout used in the sampler, correct?

@user1:
I think so.

diff_hunk: @@ -0,0 +1,388 @@
+"""Torch Module for Relational graph convolution layer using the
+aggregation primitives in cugraph-ops"""
+# pylint: disable=no-member, arguments-differ, invalid-name, too-many-arguments
+import math
+
+import torch as th
+from torch import nn
+
+try:
+    from pylibcugraphops.legacy.aggregators.node_level import (
+        agg_hg_basis_post_bwd_int32,
+        agg_hg_basis_post_bwd_int64,
+        agg_hg_basis_post_fwd_int32,
+        agg_hg_basis_post_fwd_int64,
+    )
+    from pylibcugraphops.legacy.structure.graph_types import (
+        message_flow_graph_hg_csr_int32,
+        message_flow_graph_hg_csr_int64,
+    )
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "dgl.nn.CuGraphRelGraphConv requires pylibcugraphops to be installed."
+    )
+
+CUDA_SM_PER_BLOCK = 49152
+
+
+class RelGraphConvAgg(th.autograd.Function):
+    r"""Custom autograd function for R-GCN aggregation layer that uses the
+    aggregation functions in cugraph-ops."""
+
+    @staticmethod
+    def forward(ctx, g, fanout, num_rels, edge_types, feat, coeff):
+        r"""Compute the forward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        fanout : int
+            Maximum in-degree of nodes.
+        num_rels : int
+            Number of relations.
+        edge_types : torch.Tensor
+            A 1D tensor of edge types. Note that pylibcugraphops only accepts
+            edge_types in int32, so any input of other integer types will be casted
+            into int32, thus introducing some overhead. Pass in int32 tensor for
+            best performance.
+        feat : torch.Tensor, dtype=torch.float32, requires_grad=True
+            Tensor of source node features. Shape: (num_src_nodes, in_feat).
+        coeff : torch.Tensor, dtype=torch.float32, requires_grad=True
+            The coefficient matrix used in basis-decomposition regularization.
+            Shape: (num_rels, num_bases). It should be set to ``None``
+            when ``regularizer=None``.
+
+        Returns
+        -------
+        agg_output : torch.Tensor, dtype=torch.float32
+            Aggregation output. Shape: (num_dst_nodes, num_rels * in_feat)
+            when ``regularizer=None``; Shape: (num_dst_nodes, num_bases * in_feat)
+            when ``regularizer='basis'``.
+        """
+        ctx.device = g.device
+        if g.idtype == th.int32:
+            mfg_csr_func = message_flow_graph_hg_csr_int32
+            agg_fwd_func = agg_hg_basis_post_fwd_int32
+        elif g.idtype == th.int64:
+            mfg_csr_func = message_flow_graph_hg_csr_int64
+            agg_fwd_func = agg_hg_basis_post_fwd_int64
+        else:
+            raise TypeError(
+                f"Supported ID type: torch.int32 or torch.int64, "
+                f"but got {g.idtype}"
+            )
+        ctx.idtype = g.idtype
+
+        _in_feat = feat.shape[-1]
+        indptr, indices, edge_ids = g.adj_sparse("csc")
+        # edge_ids is in a mixed order, need to permutate incoming etypes
+        # and stash the result for backward propagation
+        ctx.edge_types_int32 = edge_types[edge_ids.long()].int()
+        # node_types are not needed by the post-variant RGCN aggregation functions
+        _num_node_types = 0
+        _out_node_types = _in_node_types = None
+
+        mfg = mfg_csr_func(
+            fanout,
+            g.dstnodes(),
+            g.srcnodes(),
+            indptr,
+            indices,
+            _num_node_types,
+            num_rels,
+            _out_node_types,
+            _in_node_types,
+            ctx.edge_types_int32,
+        )
+        ctx.mfg = mfg
+
+        if coeff is None:
+            leading_dimension = num_rels * _in_feat
+        else:
+            _num_bases = coeff.shape[-1]
+            leading_dimension = _num_bases * _in_feat
+
+        agg_output = th.empty(
+            g.num_dst_nodes(),
+            leading_dimension,
+            dtype=th.float32,
+            device=ctx.device,
+        )
+        if coeff is None:
+            agg_fwd_func(agg_output, feat.detach(), mfg)
+        else:
+            agg_fwd_func(
+                agg_output,
+                feat.detach(),
+                mfg,
+                weights_combination=coeff.detach()
+            )
+
+        ctx.save_for_backward(feat, coeff)
+        return agg_output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        r"""Compute the backward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        grad_output : torch.Tensor, dtype=torch.float32
+            Gradient of loss function w.r.t output.
+        """
+        feat, coeff = ctx.saved_tensors
+
+        if ctx.idtype == th.int32:
+            agg_bwd_func = agg_hg_basis_post_bwd_int32
+        else:
+            agg_bwd_func = agg_hg_basis_post_bwd_int64
+
+        grad_feat = th.empty_like(feat, dtype=th.float32, device=ctx.device)
+        if coeff is None:
+            grad_coeff = None
+            agg_bwd_func(grad_feat, grad_output, feat.detach(), ctx.mfg)
+        else:
+            grad_coeff = th.empty_like(
+                coeff, dtype=th.float32, device=ctx.device
+            )
+            agg_bwd_func(
+                grad_feat,
+                grad_output,
+                feat.detach(),
+                ctx.mfg,
+                output_weight_gradient=grad_coeff,
+                weights_combination=coeff.detach()
+            )
+
+        return None, None, None, None, grad_feat, grad_coeff
+
+
+class CuGraphRelGraphConv(nn.Module):
+    r"""An accelerated relational graph convolution layer using the
+    highly-optimized aggregation primitives in cugraph-ops.
+    See :class:`dgl.nn.pytorch.conv.RelGraphConv` for mathematical model.
+
+    .. note::
+        This is an **experimental** feature.
+        Compared with :class:`dgl.nn.pytorch.conv.RelGraphConv`, this model:
+
+        * Only works on cuda devices.
+        * Only support basis-decomposition regularization.
+        * Requires an extra argument `fanout` as input, since the current model is
+          designed for sampled-graph (message-flow-graph) use-cases.
+          Full-graph support will be added in upcoming releases.
+
+    Parameters
+    ----------
+    in_feat : int
+        Input feature size.
+    out_feat : int
+        Output feature size.
+    num_rels : int
+        Number of relations.
+    fanout : int
+        Maximum number of sampled neighbors of an destination node;
+        i.e, maximum in degree of destination nodes
+    regularizer : str, optional
+        Which weight regularizer to use ("basis" or ``None``):
+         - "basis" is for basis-decomposition.
+         - ``None`` applies no regularization.
+        Default: ``None``.
+    num_bases : int, optional
+        Number of bases. Only used when "basis" regularizer is applied.
+        Default: ``None``.
+    bias : bool, optional
+        True if bias is added. Default: ``True``.
+    activation : callable, optional
+        Activation function. Default: ``None``.
+    self_loop : bool, optional
+        True to include self loop message. Default: ``True``.
+    dropout : float, optional
+        Dropout rate. Default: ``0.0``
+    layer_norm: bool, optional
path: python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py
text: @user1:
```suggestion
    layer_norm : bool, optional
```

@author:
Fixed.

diff_hunk: @@ -0,0 +1,388 @@
+"""Torch Module for Relational graph convolution layer using the
+aggregation primitives in cugraph-ops"""
+# pylint: disable=no-member, arguments-differ, invalid-name, too-many-arguments
+import math
+
+import torch as th
+from torch import nn
+
+try:
+    from pylibcugraphops.legacy.aggregators.node_level import (
+        agg_hg_basis_post_bwd_int32,
+        agg_hg_basis_post_bwd_int64,
+        agg_hg_basis_post_fwd_int32,
+        agg_hg_basis_post_fwd_int64,
+    )
+    from pylibcugraphops.legacy.structure.graph_types import (
+        message_flow_graph_hg_csr_int32,
+        message_flow_graph_hg_csr_int64,
+    )
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "dgl.nn.CuGraphRelGraphConv requires pylibcugraphops to be installed."
+    )
+
+CUDA_SM_PER_BLOCK = 49152
+
+
+class RelGraphConvAgg(th.autograd.Function):
+    r"""Custom autograd function for R-GCN aggregation layer that uses the
+    aggregation functions in cugraph-ops."""
+
+    @staticmethod
+    def forward(ctx, g, fanout, num_rels, edge_types, feat, coeff):
+        r"""Compute the forward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        fanout : int
+            Maximum in-degree of nodes.
+        num_rels : int
+            Number of relations.
+        edge_types : torch.Tensor
+            A 1D tensor of edge types. Note that pylibcugraphops only accepts
+            edge_types in int32, so any input of other integer types will be casted
+            into int32, thus introducing some overhead. Pass in int32 tensor for
+            best performance.
+        feat : torch.Tensor, dtype=torch.float32, requires_grad=True
+            Tensor of source node features. Shape: (num_src_nodes, in_feat).
+        coeff : torch.Tensor, dtype=torch.float32, requires_grad=True
+            The coefficient matrix used in basis-decomposition regularization.
+            Shape: (num_rels, num_bases). It should be set to ``None``
+            when ``regularizer=None``.
+
+        Returns
+        -------
+        agg_output : torch.Tensor, dtype=torch.float32
+            Aggregation output. Shape: (num_dst_nodes, num_rels * in_feat)
+            when ``regularizer=None``; Shape: (num_dst_nodes, num_bases * in_feat)
+            when ``regularizer='basis'``.
+        """
+        ctx.device = g.device
+        if g.idtype == th.int32:
+            mfg_csr_func = message_flow_graph_hg_csr_int32
+            agg_fwd_func = agg_hg_basis_post_fwd_int32
+        elif g.idtype == th.int64:
+            mfg_csr_func = message_flow_graph_hg_csr_int64
+            agg_fwd_func = agg_hg_basis_post_fwd_int64
+        else:
+            raise TypeError(
+                f"Supported ID type: torch.int32 or torch.int64, "
+                f"but got {g.idtype}"
+            )
+        ctx.idtype = g.idtype
+
+        _in_feat = feat.shape[-1]
+        indptr, indices, edge_ids = g.adj_sparse("csc")
+        # edge_ids is in a mixed order, need to permutate incoming etypes
+        # and stash the result for backward propagation
+        ctx.edge_types_int32 = edge_types[edge_ids.long()].int()
+        # node_types are not needed by the post-variant RGCN aggregation functions
+        _num_node_types = 0
+        _out_node_types = _in_node_types = None
+
+        mfg = mfg_csr_func(
+            fanout,
+            g.dstnodes(),
+            g.srcnodes(),
+            indptr,
+            indices,
+            _num_node_types,
+            num_rels,
+            _out_node_types,
+            _in_node_types,
+            ctx.edge_types_int32,
+        )
+        ctx.mfg = mfg
+
+        if coeff is None:
+            leading_dimension = num_rels * _in_feat
+        else:
+            _num_bases = coeff.shape[-1]
+            leading_dimension = _num_bases * _in_feat
+
+        agg_output = th.empty(
+            g.num_dst_nodes(),
+            leading_dimension,
+            dtype=th.float32,
+            device=ctx.device,
+        )
+        if coeff is None:
+            agg_fwd_func(agg_output, feat.detach(), mfg)
+        else:
+            agg_fwd_func(
+                agg_output,
+                feat.detach(),
+                mfg,
+                weights_combination=coeff.detach()
+            )
+
+        ctx.save_for_backward(feat, coeff)
+        return agg_output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        r"""Compute the backward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        grad_output : torch.Tensor, dtype=torch.float32
+            Gradient of loss function w.r.t output.
+        """
+        feat, coeff = ctx.saved_tensors
+
+        if ctx.idtype == th.int32:
+            agg_bwd_func = agg_hg_basis_post_bwd_int32
+        else:
+            agg_bwd_func = agg_hg_basis_post_bwd_int64
+
+        grad_feat = th.empty_like(feat, dtype=th.float32, device=ctx.device)
+        if coeff is None:
+            grad_coeff = None
+            agg_bwd_func(grad_feat, grad_output, feat.detach(), ctx.mfg)
+        else:
+            grad_coeff = th.empty_like(
+                coeff, dtype=th.float32, device=ctx.device
+            )
+            agg_bwd_func(
+                grad_feat,
+                grad_output,
+                feat.detach(),
+                ctx.mfg,
+                output_weight_gradient=grad_coeff,
+                weights_combination=coeff.detach()
+            )
+
+        return None, None, None, None, grad_feat, grad_coeff
+
+
+class CuGraphRelGraphConv(nn.Module):
+    r"""An accelerated relational graph convolution layer using the
+    highly-optimized aggregation primitives in cugraph-ops.
+    See :class:`dgl.nn.pytorch.conv.RelGraphConv` for mathematical model.
+
+    .. note::
+        This is an **experimental** feature.
+        Compared with :class:`dgl.nn.pytorch.conv.RelGraphConv`, this model:
+
+        * Only works on cuda devices.
+        * Only support basis-decomposition regularization.
+        * Requires an extra argument `fanout` as input, since the current model is
+          designed for sampled-graph (message-flow-graph) use-cases.
+          Full-graph support will be added in upcoming releases.
+
+    Parameters
+    ----------
+    in_feat : int
+        Input feature size.
+    out_feat : int
+        Output feature size.
+    num_rels : int
+        Number of relations.
+    fanout : int
+        Maximum number of sampled neighbors of an destination node;
+        i.e, maximum in degree of destination nodes
+    regularizer : str, optional
+        Which weight regularizer to use ("basis" or ``None``):
+         - "basis" is for basis-decomposition.
+         - ``None`` applies no regularization.
+        Default: ``None``.
+    num_bases : int, optional
+        Number of bases. Only used when "basis" regularizer is applied.
path: python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py
text: @user1:
```suggestion
        Number of bases. It comes into effect when the "basis" regularizer is applied. Otherwise, it uses number of relations (``num_rels``).
```

@author:
> Otherwise, it uses number of relations (``num_rels``).

This is not the case in the current setup. The value is not being used when regularizer is None. I did add "It comes into effect when the "basis" regularizer is applied." for consistency with RelGraphConv.

@user1:
I see. Thanks.

diff_hunk: @@ -0,0 +1,388 @@
+"""Torch Module for Relational graph convolution layer using the
+aggregation primitives in cugraph-ops"""
+# pylint: disable=no-member, arguments-differ, invalid-name, too-many-arguments
+import math
+
+import torch as th
+from torch import nn
+
+try:
+    from pylibcugraphops.legacy.aggregators.node_level import (
+        agg_hg_basis_post_bwd_int32,
+        agg_hg_basis_post_bwd_int64,
+        agg_hg_basis_post_fwd_int32,
+        agg_hg_basis_post_fwd_int64,
+    )
+    from pylibcugraphops.legacy.structure.graph_types import (
+        message_flow_graph_hg_csr_int32,
+        message_flow_graph_hg_csr_int64,
+    )
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "dgl.nn.CuGraphRelGraphConv requires pylibcugraphops to be installed."
+    )
+
+CUDA_SM_PER_BLOCK = 49152
+
+
+class RelGraphConvAgg(th.autograd.Function):
+    r"""Custom autograd function for R-GCN aggregation layer that uses the
+    aggregation functions in cugraph-ops."""
+
+    @staticmethod
+    def forward(ctx, g, fanout, num_rels, edge_types, feat, coeff):
+        r"""Compute the forward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        fanout : int
+            Maximum in-degree of nodes.
+        num_rels : int
+            Number of relations.
+        edge_types : torch.Tensor
+            A 1D tensor of edge types. Note that pylibcugraphops only accepts
+            edge_types in int32, so any input of other integer types will be casted
+            into int32, thus introducing some overhead. Pass in int32 tensor for
+            best performance.
+        feat : torch.Tensor, dtype=torch.float32, requires_grad=True
+            Tensor of source node features. Shape: (num_src_nodes, in_feat).
+        coeff : torch.Tensor, dtype=torch.float32, requires_grad=True
+            The coefficient matrix used in basis-decomposition regularization.
+            Shape: (num_rels, num_bases). It should be set to ``None``
+            when ``regularizer=None``.
+
+        Returns
+        -------
+        agg_output : torch.Tensor, dtype=torch.float32
+            Aggregation output. Shape: (num_dst_nodes, num_rels * in_feat)
+            when ``regularizer=None``; Shape: (num_dst_nodes, num_bases * in_feat)
+            when ``regularizer='basis'``.
+        """
+        ctx.device = g.device
+        if g.idtype == th.int32:
+            mfg_csr_func = message_flow_graph_hg_csr_int32
+            agg_fwd_func = agg_hg_basis_post_fwd_int32
+        elif g.idtype == th.int64:
+            mfg_csr_func = message_flow_graph_hg_csr_int64
+            agg_fwd_func = agg_hg_basis_post_fwd_int64
+        else:
+            raise TypeError(
+                f"Supported ID type: torch.int32 or torch.int64, "
+                f"but got {g.idtype}"
+            )
+        ctx.idtype = g.idtype
+
+        _in_feat = feat.shape[-1]
+        indptr, indices, edge_ids = g.adj_sparse("csc")
+        # edge_ids is in a mixed order, need to permutate incoming etypes
+        # and stash the result for backward propagation
+        ctx.edge_types_int32 = edge_types[edge_ids.long()].int()
+        # node_types are not needed by the post-variant RGCN aggregation functions
+        _num_node_types = 0
+        _out_node_types = _in_node_types = None
+
+        mfg = mfg_csr_func(
+            fanout,
+            g.dstnodes(),
+            g.srcnodes(),
+            indptr,
+            indices,
+            _num_node_types,
+            num_rels,
+            _out_node_types,
+            _in_node_types,
+            ctx.edge_types_int32,
+        )
+        ctx.mfg = mfg
+
+        if coeff is None:
+            leading_dimension = num_rels * _in_feat
+        else:
+            _num_bases = coeff.shape[-1]
+            leading_dimension = _num_bases * _in_feat
+
+        agg_output = th.empty(
+            g.num_dst_nodes(),
+            leading_dimension,
+            dtype=th.float32,
+            device=ctx.device,
+        )
+        if coeff is None:
+            agg_fwd_func(agg_output, feat.detach(), mfg)
+        else:
+            agg_fwd_func(
+                agg_output,
+                feat.detach(),
+                mfg,
+                weights_combination=coeff.detach()
+            )
+
+        ctx.save_for_backward(feat, coeff)
+        return agg_output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        r"""Compute the backward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        grad_output : torch.Tensor, dtype=torch.float32
+            Gradient of loss function w.r.t output.
+        """
+        feat, coeff = ctx.saved_tensors
+
+        if ctx.idtype == th.int32:
+            agg_bwd_func = agg_hg_basis_post_bwd_int32
+        else:
+            agg_bwd_func = agg_hg_basis_post_bwd_int64
+
+        grad_feat = th.empty_like(feat, dtype=th.float32, device=ctx.device)
+        if coeff is None:
+            grad_coeff = None
+            agg_bwd_func(grad_feat, grad_output, feat.detach(), ctx.mfg)
+        else:
+            grad_coeff = th.empty_like(
+                coeff, dtype=th.float32, device=ctx.device
+            )
+            agg_bwd_func(
+                grad_feat,
+                grad_output,
+                feat.detach(),
+                ctx.mfg,
+                output_weight_gradient=grad_coeff,
+                weights_combination=coeff.detach()
+            )
+
+        return None, None, None, None, grad_feat, grad_coeff
+
+
+class CuGraphRelGraphConv(nn.Module):
path: python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py
text: @user1:
Any instructions on installing the dependencies? I remember CuGraph only support a subset of CUDA versions. Is there a easy recipe for installing all these dependencies?

@author:
That's a good point. Unlike what its name suggests, this module does not depend on cugraph. It only requires the `pylibcugraphops` package, which can be installed via `conda install -c nvidia pylibcugraphops=22.12` after its release.
It works with any CUDA toolkit version >= 11.

Where do you suggest to add this information?

@user1:
For now, you can simply add these two requirements in the doc of `CuGraphRelGraphConv`.

@author:
Instruction added.

diff_hunk: @@ -0,0 +1,388 @@
+"""Torch Module for Relational graph convolution layer using the
+aggregation primitives in cugraph-ops"""
+# pylint: disable=no-member, arguments-differ, invalid-name, too-many-arguments
+import math
+
+import torch as th
+from torch import nn
+
+try:
+    from pylibcugraphops.legacy.aggregators.node_level import (
+        agg_hg_basis_post_bwd_int32,
+        agg_hg_basis_post_bwd_int64,
+        agg_hg_basis_post_fwd_int32,
+        agg_hg_basis_post_fwd_int64,
+    )
+    from pylibcugraphops.legacy.structure.graph_types import (
+        message_flow_graph_hg_csr_int32,
+        message_flow_graph_hg_csr_int64,
+    )
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "dgl.nn.CuGraphRelGraphConv requires pylibcugraphops to be installed."
+    )
+
+CUDA_SM_PER_BLOCK = 49152
+
+
+class RelGraphConvAgg(th.autograd.Function):
+    r"""Custom autograd function for R-GCN aggregation layer that uses the
+    aggregation functions in cugraph-ops."""
+
+    @staticmethod
+    def forward(ctx, g, fanout, num_rels, edge_types, feat, coeff):
+        r"""Compute the forward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        fanout : int
+            Maximum in-degree of nodes.
+        num_rels : int
+            Number of relations.
+        edge_types : torch.Tensor
+            A 1D tensor of edge types. Note that pylibcugraphops only accepts
+            edge_types in int32, so any input of other integer types will be casted
+            into int32, thus introducing some overhead. Pass in int32 tensor for
+            best performance.
+        feat : torch.Tensor, dtype=torch.float32, requires_grad=True
+            Tensor of source node features. Shape: (num_src_nodes, in_feat).
+        coeff : torch.Tensor, dtype=torch.float32, requires_grad=True
+            The coefficient matrix used in basis-decomposition regularization.
+            Shape: (num_rels, num_bases). It should be set to ``None``
+            when ``regularizer=None``.
+
+        Returns
+        -------
+        agg_output : torch.Tensor, dtype=torch.float32
+            Aggregation output. Shape: (num_dst_nodes, num_rels * in_feat)
+            when ``regularizer=None``; Shape: (num_dst_nodes, num_bases * in_feat)
+            when ``regularizer='basis'``.
+        """
+        ctx.device = g.device
+        if g.idtype == th.int32:
+            mfg_csr_func = message_flow_graph_hg_csr_int32
+            agg_fwd_func = agg_hg_basis_post_fwd_int32
+        elif g.idtype == th.int64:
+            mfg_csr_func = message_flow_graph_hg_csr_int64
+            agg_fwd_func = agg_hg_basis_post_fwd_int64
+        else:
+            raise TypeError(
+                f"Supported ID type: torch.int32 or torch.int64, "
+                f"but got {g.idtype}"
+            )
+        ctx.idtype = g.idtype
+
+        _in_feat = feat.shape[-1]
+        indptr, indices, edge_ids = g.adj_sparse("csc")
+        # edge_ids is in a mixed order, need to permutate incoming etypes
+        # and stash the result for backward propagation
+        ctx.edge_types_int32 = edge_types[edge_ids.long()].int()
+        # node_types are not needed by the post-variant RGCN aggregation functions
+        _num_node_types = 0
+        _out_node_types = _in_node_types = None
+
+        mfg = mfg_csr_func(
+            fanout,
+            g.dstnodes(),
+            g.srcnodes(),
+            indptr,
+            indices,
+            _num_node_types,
+            num_rels,
+            _out_node_types,
+            _in_node_types,
+            ctx.edge_types_int32,
+        )
+        ctx.mfg = mfg
+
+        if coeff is None:
+            leading_dimension = num_rels * _in_feat
+        else:
+            _num_bases = coeff.shape[-1]
+            leading_dimension = _num_bases * _in_feat
+
+        agg_output = th.empty(
+            g.num_dst_nodes(),
+            leading_dimension,
+            dtype=th.float32,
+            device=ctx.device,
+        )
+        if coeff is None:
+            agg_fwd_func(agg_output, feat.detach(), mfg)
+        else:
+            agg_fwd_func(
+                agg_output,
+                feat.detach(),
+                mfg,
+                weights_combination=coeff.detach()
+            )
+
+        ctx.save_for_backward(feat, coeff)
+        return agg_output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        r"""Compute the backward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        grad_output : torch.Tensor, dtype=torch.float32
+            Gradient of loss function w.r.t output.
+        """
+        feat, coeff = ctx.saved_tensors
+
+        if ctx.idtype == th.int32:
+            agg_bwd_func = agg_hg_basis_post_bwd_int32
+        else:
+            agg_bwd_func = agg_hg_basis_post_bwd_int64
+
+        grad_feat = th.empty_like(feat, dtype=th.float32, device=ctx.device)
+        if coeff is None:
+            grad_coeff = None
+            agg_bwd_func(grad_feat, grad_output, feat.detach(), ctx.mfg)
+        else:
+            grad_coeff = th.empty_like(
+                coeff, dtype=th.float32, device=ctx.device
+            )
+            agg_bwd_func(
+                grad_feat,
+                grad_output,
+                feat.detach(),
+                ctx.mfg,
+                output_weight_gradient=grad_coeff,
+                weights_combination=coeff.detach()
+            )
+
+        return None, None, None, None, grad_feat, grad_coeff
+
+
+class CuGraphRelGraphConv(nn.Module):
+    r"""An accelerated relational graph convolution layer using the
+    highly-optimized aggregation primitives in cugraph-ops.
+    See :class:`dgl.nn.pytorch.conv.RelGraphConv` for mathematical model.
+
+    .. note::
+        This is an **experimental** feature.
+        Compared with :class:`dgl.nn.pytorch.conv.RelGraphConv`, this model:
+
+        * Only works on cuda devices.
+        * Only support basis-decomposition regularization.
+        * Requires an extra argument `fanout` as input, since the current model is
+          designed for sampled-graph (message-flow-graph) use-cases.
+          Full-graph support will be added in upcoming releases.
+
+    Parameters
+    ----------
+    in_feat : int
+        Input feature size.
+    out_feat : int
+        Output feature size.
+    num_rels : int
+        Number of relations.
+    fanout : int
+        Maximum number of sampled neighbors of an destination node;
+        i.e, maximum in degree of destination nodes
+    regularizer : str, optional
+        Which weight regularizer to use ("basis" or ``None``):
+         - "basis" is for basis-decomposition.
+         - ``None`` applies no regularization.
+        Default: ``None``.
+    num_bases : int, optional
+        Number of bases. Only used when "basis" regularizer is applied.
+        Default: ``None``.
+    bias : bool, optional
+        True if bias is added. Default: ``True``.
+    activation : callable, optional
+        Activation function. Default: ``None``.
+    self_loop : bool, optional
+        True to include self loop message. Default: ``True``.
+    dropout : float, optional
+        Dropout rate. Default: ``0.0``
+    layer_norm: bool, optional
+        True to add layer norm. Default: ``False``
+
+    Examples
+    --------
+    >>> import dgl
+    >>> from dgl.dataloading import NeighborSampler, DataLoader
+    >>> from dgl.nn import CuGraphRelGraphConv
+    >>> import torch.nn.functional as F
+    ...
+    >>> device = 'cuda'
+    >>> fanouts = [5, 6]
+    >>> sampler = NeighborSampler(fanouts)
+    >>> dataloader = DataLoader(g, train_nid, sampler,
+    ...     device=device, batch_size=1024)
+    >>> conv1 = CuGraphRelGraphConv(in_dim, h_dim, num_rels, fanouts[0],
+    ...     regularizer='basis', num_bases=10).to(device)
+    >>> conv2 = CuGraphRelGraphConv(h_dim, out_dim, num_rels, fanouts[1],
+    ...     regularizer='basis', num_bases=10).to(device)
+    >>> for input_nodes, output_nodes, blocks in dataloader:
+    ...     h = conv1(blocks[0], x, blocks[0].edata[dgl.ETYPE])
+    ...     h = F.relu(h)
+    ...     h = conv2(blocks[1], h, blocks[1].edata[dgl.ETYPE])
+    """
+    def __init__(
+        self,
+        in_feat,
+        out_feat,
+        num_rels,
+        fanout,
+        regularizer=None,
+        num_bases=None,
+        bias=True,
+        activation=None,
+        self_loop=True,
+        dropout=0.0,
+        layer_norm=False,
+    ):
+        super().__init__()
+        self.in_feat = in_feat
+        self.out_feat = out_feat
+        self.num_rels = num_rels
+        self.fanout = fanout
+
+        # regularizer (see dgl.nn.pytorch.linear.TypedLinear)
+        if regularizer is None:
+            self.W = nn.Parameter(th.Tensor(num_rels, in_feat, out_feat))
+            self.coeff = None
+        elif regularizer == "basis":
+            if num_bases is None:
+                raise ValueError(
+                    'Missing "num_bases" for basis regularization.'
+                )
+            self.W = nn.Parameter(th.Tensor(num_bases, in_feat, out_feat))
+            self.coeff = nn.Parameter(th.Tensor(num_rels, num_bases))
+            self.num_bases = num_bases
+        else:
+            raise ValueError(
+                f"Supported regularizer options: 'basis', "
+                f"but got {regularizer}"
+            )
+        self.regularizer = regularizer
+        self.reset_parameters()
+
+        # others
+        self.bias = bias
+        self.activation = activation
+        self.self_loop = self_loop
+        self.layer_norm = layer_norm
+
+        # bias
+        if self.bias:
+            self.h_bias = nn.Parameter(th.Tensor(out_feat))
+            nn.init.zeros_(self.h_bias)
+
+        # layer norm
+        if self.layer_norm:
+            self.layer_norm_weight = nn.LayerNorm(
+                out_feat, elementwise_affine=True
+            )
+
+        # weight for self_loop
+        if self.self_loop:
+            self.loop_weight = nn.Parameter(th.Tensor(in_feat, out_feat))
+            nn.init.xavier_uniform_(
+                self.loop_weight, gain=nn.init.calculate_gain("relu")
+            )
+
+        self.dropout = nn.Dropout(dropout)
+
+    def reset_parameters(self):
path: python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py
text: @user1:
Typically `reset_parameters` allows users to reset all learnable parameters while this currently only handles `self.coeff` and `self.W`. 

- We can either do these directly in `__init__` if this is only used by `__init__`
- Or expand it to cover all learnable parameters

@author:
great suggestion. The initialization calls now go into __init__.

diff_hunk: @@ -0,0 +1,388 @@
+"""Torch Module for Relational graph convolution layer using the
+aggregation primitives in cugraph-ops"""
+# pylint: disable=no-member, arguments-differ, invalid-name, too-many-arguments
+import math
+
+import torch as th
+from torch import nn
+
+try:
+    from pylibcugraphops.legacy.aggregators.node_level import (
+        agg_hg_basis_post_bwd_int32,
+        agg_hg_basis_post_bwd_int64,
+        agg_hg_basis_post_fwd_int32,
+        agg_hg_basis_post_fwd_int64,
+    )
+    from pylibcugraphops.legacy.structure.graph_types import (
+        message_flow_graph_hg_csr_int32,
+        message_flow_graph_hg_csr_int64,
+    )
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "dgl.nn.CuGraphRelGraphConv requires pylibcugraphops to be installed."
+    )
+
+CUDA_SM_PER_BLOCK = 49152
+
+
+class RelGraphConvAgg(th.autograd.Function):
+    r"""Custom autograd function for R-GCN aggregation layer that uses the
+    aggregation functions in cugraph-ops."""
+
+    @staticmethod
+    def forward(ctx, g, fanout, num_rels, edge_types, feat, coeff):
+        r"""Compute the forward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        fanout : int
+            Maximum in-degree of nodes.
+        num_rels : int
+            Number of relations.
+        edge_types : torch.Tensor
+            A 1D tensor of edge types. Note that pylibcugraphops only accepts
+            edge_types in int32, so any input of other integer types will be casted
+            into int32, thus introducing some overhead. Pass in int32 tensor for
+            best performance.
+        feat : torch.Tensor, dtype=torch.float32, requires_grad=True
+            Tensor of source node features. Shape: (num_src_nodes, in_feat).
+        coeff : torch.Tensor, dtype=torch.float32, requires_grad=True
+            The coefficient matrix used in basis-decomposition regularization.
+            Shape: (num_rels, num_bases). It should be set to ``None``
+            when ``regularizer=None``.
+
+        Returns
+        -------
+        agg_output : torch.Tensor, dtype=torch.float32
+            Aggregation output. Shape: (num_dst_nodes, num_rels * in_feat)
+            when ``regularizer=None``; Shape: (num_dst_nodes, num_bases * in_feat)
+            when ``regularizer='basis'``.
+        """
+        ctx.device = g.device
+        if g.idtype == th.int32:
+            mfg_csr_func = message_flow_graph_hg_csr_int32
+            agg_fwd_func = agg_hg_basis_post_fwd_int32
+        elif g.idtype == th.int64:
+            mfg_csr_func = message_flow_graph_hg_csr_int64
+            agg_fwd_func = agg_hg_basis_post_fwd_int64
+        else:
+            raise TypeError(
+                f"Supported ID type: torch.int32 or torch.int64, "
+                f"but got {g.idtype}"
+            )
+        ctx.idtype = g.idtype
+
+        _in_feat = feat.shape[-1]
+        indptr, indices, edge_ids = g.adj_sparse("csc")
+        # edge_ids is in a mixed order, need to permutate incoming etypes
+        # and stash the result for backward propagation
+        ctx.edge_types_int32 = edge_types[edge_ids.long()].int()
+        # node_types are not needed by the post-variant RGCN aggregation functions
+        _num_node_types = 0
+        _out_node_types = _in_node_types = None
+
+        mfg = mfg_csr_func(
+            fanout,
+            g.dstnodes(),
+            g.srcnodes(),
+            indptr,
+            indices,
+            _num_node_types,
+            num_rels,
+            _out_node_types,
+            _in_node_types,
+            ctx.edge_types_int32,
+        )
+        ctx.mfg = mfg
+
+        if coeff is None:
+            leading_dimension = num_rels * _in_feat
+        else:
+            _num_bases = coeff.shape[-1]
+            leading_dimension = _num_bases * _in_feat
+
+        agg_output = th.empty(
+            g.num_dst_nodes(),
+            leading_dimension,
+            dtype=th.float32,
+            device=ctx.device,
+        )
+        if coeff is None:
+            agg_fwd_func(agg_output, feat.detach(), mfg)
+        else:
+            agg_fwd_func(
+                agg_output,
+                feat.detach(),
+                mfg,
+                weights_combination=coeff.detach()
+            )
+
+        ctx.save_for_backward(feat, coeff)
+        return agg_output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        r"""Compute the backward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        grad_output : torch.Tensor, dtype=torch.float32
+            Gradient of loss function w.r.t output.
+        """
+        feat, coeff = ctx.saved_tensors
+
+        if ctx.idtype == th.int32:
+            agg_bwd_func = agg_hg_basis_post_bwd_int32
+        else:
+            agg_bwd_func = agg_hg_basis_post_bwd_int64
+
+        grad_feat = th.empty_like(feat, dtype=th.float32, device=ctx.device)
+        if coeff is None:
+            grad_coeff = None
+            agg_bwd_func(grad_feat, grad_output, feat.detach(), ctx.mfg)
+        else:
+            grad_coeff = th.empty_like(
+                coeff, dtype=th.float32, device=ctx.device
+            )
+            agg_bwd_func(
+                grad_feat,
+                grad_output,
+                feat.detach(),
+                ctx.mfg,
+                output_weight_gradient=grad_coeff,
+                weights_combination=coeff.detach()
+            )
+
+        return None, None, None, None, grad_feat, grad_coeff
+
+
+class CuGraphRelGraphConv(nn.Module):
+    r"""An accelerated relational graph convolution layer using the
+    highly-optimized aggregation primitives in cugraph-ops.
+    See :class:`dgl.nn.pytorch.conv.RelGraphConv` for mathematical model.
+
+    .. note::
+        This is an **experimental** feature.
+        Compared with :class:`dgl.nn.pytorch.conv.RelGraphConv`, this model:
+
+        * Only works on cuda devices.
+        * Only support basis-decomposition regularization.
+        * Requires an extra argument `fanout` as input, since the current model is
+          designed for sampled-graph (message-flow-graph) use-cases.
+          Full-graph support will be added in upcoming releases.
+
+    Parameters
+    ----------
+    in_feat : int
+        Input feature size.
+    out_feat : int
+        Output feature size.
+    num_rels : int
+        Number of relations.
+    fanout : int
+        Maximum number of sampled neighbors of an destination node;
+        i.e, maximum in degree of destination nodes
+    regularizer : str, optional
+        Which weight regularizer to use ("basis" or ``None``):
+         - "basis" is for basis-decomposition.
+         - ``None`` applies no regularization.
+        Default: ``None``.
+    num_bases : int, optional
+        Number of bases. Only used when "basis" regularizer is applied.
+        Default: ``None``.
+    bias : bool, optional
+        True if bias is added. Default: ``True``.
+    activation : callable, optional
+        Activation function. Default: ``None``.
+    self_loop : bool, optional
+        True to include self loop message. Default: ``True``.
+    dropout : float, optional
+        Dropout rate. Default: ``0.0``
+    layer_norm: bool, optional
+        True to add layer norm. Default: ``False``
+
+    Examples
+    --------
+    >>> import dgl
+    >>> from dgl.dataloading import NeighborSampler, DataLoader
+    >>> from dgl.nn import CuGraphRelGraphConv
+    >>> import torch.nn.functional as F
+    ...
+    >>> device = 'cuda'
+    >>> fanouts = [5, 6]
+    >>> sampler = NeighborSampler(fanouts)
+    >>> dataloader = DataLoader(g, train_nid, sampler,
+    ...     device=device, batch_size=1024)
+    >>> conv1 = CuGraphRelGraphConv(in_dim, h_dim, num_rels, fanouts[0],
+    ...     regularizer='basis', num_bases=10).to(device)
+    >>> conv2 = CuGraphRelGraphConv(h_dim, out_dim, num_rels, fanouts[1],
+    ...     regularizer='basis', num_bases=10).to(device)
+    >>> for input_nodes, output_nodes, blocks in dataloader:
+    ...     h = conv1(blocks[0], x, blocks[0].edata[dgl.ETYPE])
+    ...     h = F.relu(h)
+    ...     h = conv2(blocks[1], h, blocks[1].edata[dgl.ETYPE])
+    """
+    def __init__(
+        self,
+        in_feat,
+        out_feat,
+        num_rels,
+        fanout,
+        regularizer=None,
+        num_bases=None,
+        bias=True,
+        activation=None,
+        self_loop=True,
+        dropout=0.0,
+        layer_norm=False,
+    ):
+        super().__init__()
+        self.in_feat = in_feat
+        self.out_feat = out_feat
+        self.num_rels = num_rels
+        self.fanout = fanout
+
+        # regularizer (see dgl.nn.pytorch.linear.TypedLinear)
+        if regularizer is None:
+            self.W = nn.Parameter(th.Tensor(num_rels, in_feat, out_feat))
+            self.coeff = None
+        elif regularizer == "basis":
+            if num_bases is None:
+                raise ValueError(
+                    'Missing "num_bases" for basis regularization.'
+                )
+            self.W = nn.Parameter(th.Tensor(num_bases, in_feat, out_feat))
+            self.coeff = nn.Parameter(th.Tensor(num_rels, num_bases))
+            self.num_bases = num_bases
+        else:
+            raise ValueError(
+                f"Supported regularizer options: 'basis', "
+                f"but got {regularizer}"
+            )
+        self.regularizer = regularizer
+        self.reset_parameters()
+
+        # others
+        self.bias = bias
+        self.activation = activation
+        self.self_loop = self_loop
+        self.layer_norm = layer_norm
+
+        # bias
+        if self.bias:
+            self.h_bias = nn.Parameter(th.Tensor(out_feat))
+            nn.init.zeros_(self.h_bias)
+
+        # layer norm
+        if self.layer_norm:
+            self.layer_norm_weight = nn.LayerNorm(
+                out_feat, elementwise_affine=True
+            )
+
+        # weight for self_loop
+        if self.self_loop:
+            self.loop_weight = nn.Parameter(th.Tensor(in_feat, out_feat))
+            nn.init.xavier_uniform_(
+                self.loop_weight, gain=nn.init.calculate_gain("relu")
+            )
+
+        self.dropout = nn.Dropout(dropout)
+
+    def reset_parameters(self):
+        with th.no_grad():
+            if self.regularizer is None:
+                nn.init.uniform_(
+                    self.W,
+                    -1 / math.sqrt(self.in_feat),
+                    1 / math.sqrt(self.in_feat),
+                )
+            elif self.regularizer == "basis":
+                nn.init.uniform_(
+                    self.W,
+                    -1 / math.sqrt(self.in_feat),
+                    1 / math.sqrt(self.in_feat),
+                )
+                nn.init.xavier_uniform_(
+                    self.coeff, gain=nn.init.calculate_gain("relu")
+                )
+            else:
+                raise ValueError(
+                    f"Supported regularizer options: 'basis', "
+                    f"but got {self.regularizer}"
+                )
+
+    def forward(self, g, feat, etypes, norm=None):
+        r"""Forward computation.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        feat : torch.Tensor
+            A 2D tensor of node features. Shape: :math:`(|V|, D_{in})`.
+        etypes : torch.Tensor or list[int]
path: python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py
text: @user1:
Only `torch.Tensor` is supported.

diff_hunk: @@ -0,0 +1,388 @@
+"""Torch Module for Relational graph convolution layer using the
+aggregation primitives in cugraph-ops"""
+# pylint: disable=no-member, arguments-differ, invalid-name, too-many-arguments
+import math
+
+import torch as th
+from torch import nn
+
+try:
+    from pylibcugraphops.legacy.aggregators.node_level import (
+        agg_hg_basis_post_bwd_int32,
+        agg_hg_basis_post_bwd_int64,
+        agg_hg_basis_post_fwd_int32,
+        agg_hg_basis_post_fwd_int64,
+    )
+    from pylibcugraphops.legacy.structure.graph_types import (
+        message_flow_graph_hg_csr_int32,
+        message_flow_graph_hg_csr_int64,
+    )
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "dgl.nn.CuGraphRelGraphConv requires pylibcugraphops to be installed."
+    )
+
+CUDA_SM_PER_BLOCK = 49152
+
+
+class RelGraphConvAgg(th.autograd.Function):
+    r"""Custom autograd function for R-GCN aggregation layer that uses the
+    aggregation functions in cugraph-ops."""
+
+    @staticmethod
+    def forward(ctx, g, fanout, num_rels, edge_types, feat, coeff):
+        r"""Compute the forward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        fanout : int
+            Maximum in-degree of nodes.
+        num_rels : int
+            Number of relations.
+        edge_types : torch.Tensor
+            A 1D tensor of edge types. Note that pylibcugraphops only accepts
+            edge_types in int32, so any input of other integer types will be casted
+            into int32, thus introducing some overhead. Pass in int32 tensor for
+            best performance.
+        feat : torch.Tensor, dtype=torch.float32, requires_grad=True
+            Tensor of source node features. Shape: (num_src_nodes, in_feat).
+        coeff : torch.Tensor, dtype=torch.float32, requires_grad=True
+            The coefficient matrix used in basis-decomposition regularization.
+            Shape: (num_rels, num_bases). It should be set to ``None``
+            when ``regularizer=None``.
+
+        Returns
+        -------
+        agg_output : torch.Tensor, dtype=torch.float32
+            Aggregation output. Shape: (num_dst_nodes, num_rels * in_feat)
+            when ``regularizer=None``; Shape: (num_dst_nodes, num_bases * in_feat)
+            when ``regularizer='basis'``.
+        """
+        ctx.device = g.device
+        if g.idtype == th.int32:
+            mfg_csr_func = message_flow_graph_hg_csr_int32
+            agg_fwd_func = agg_hg_basis_post_fwd_int32
+        elif g.idtype == th.int64:
+            mfg_csr_func = message_flow_graph_hg_csr_int64
+            agg_fwd_func = agg_hg_basis_post_fwd_int64
+        else:
+            raise TypeError(
+                f"Supported ID type: torch.int32 or torch.int64, "
+                f"but got {g.idtype}"
+            )
+        ctx.idtype = g.idtype
+
+        _in_feat = feat.shape[-1]
+        indptr, indices, edge_ids = g.adj_sparse("csc")
+        # edge_ids is in a mixed order, need to permutate incoming etypes
+        # and stash the result for backward propagation
+        ctx.edge_types_int32 = edge_types[edge_ids.long()].int()
+        # node_types are not needed by the post-variant RGCN aggregation functions
+        _num_node_types = 0
+        _out_node_types = _in_node_types = None
+
+        mfg = mfg_csr_func(
+            fanout,
+            g.dstnodes(),
+            g.srcnodes(),
+            indptr,
+            indices,
+            _num_node_types,
+            num_rels,
+            _out_node_types,
+            _in_node_types,
+            ctx.edge_types_int32,
+        )
+        ctx.mfg = mfg
+
+        if coeff is None:
+            leading_dimension = num_rels * _in_feat
+        else:
+            _num_bases = coeff.shape[-1]
+            leading_dimension = _num_bases * _in_feat
+
+        agg_output = th.empty(
+            g.num_dst_nodes(),
+            leading_dimension,
+            dtype=th.float32,
+            device=ctx.device,
+        )
+        if coeff is None:
+            agg_fwd_func(agg_output, feat.detach(), mfg)
+        else:
+            agg_fwd_func(
+                agg_output,
+                feat.detach(),
+                mfg,
+                weights_combination=coeff.detach()
+            )
+
+        ctx.save_for_backward(feat, coeff)
+        return agg_output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        r"""Compute the backward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        grad_output : torch.Tensor, dtype=torch.float32
+            Gradient of loss function w.r.t output.
+        """
+        feat, coeff = ctx.saved_tensors
+
+        if ctx.idtype == th.int32:
+            agg_bwd_func = agg_hg_basis_post_bwd_int32
+        else:
+            agg_bwd_func = agg_hg_basis_post_bwd_int64
+
+        grad_feat = th.empty_like(feat, dtype=th.float32, device=ctx.device)
+        if coeff is None:
+            grad_coeff = None
+            agg_bwd_func(grad_feat, grad_output, feat.detach(), ctx.mfg)
+        else:
+            grad_coeff = th.empty_like(
+                coeff, dtype=th.float32, device=ctx.device
+            )
+            agg_bwd_func(
+                grad_feat,
+                grad_output,
+                feat.detach(),
+                ctx.mfg,
+                output_weight_gradient=grad_coeff,
+                weights_combination=coeff.detach()
+            )
+
+        return None, None, None, None, grad_feat, grad_coeff
+
+
+class CuGraphRelGraphConv(nn.Module):
+    r"""An accelerated relational graph convolution layer using the
+    highly-optimized aggregation primitives in cugraph-ops.
+    See :class:`dgl.nn.pytorch.conv.RelGraphConv` for mathematical model.
+
+    .. note::
+        This is an **experimental** feature.
+        Compared with :class:`dgl.nn.pytorch.conv.RelGraphConv`, this model:
+
+        * Only works on cuda devices.
+        * Only support basis-decomposition regularization.
+        * Requires an extra argument `fanout` as input, since the current model is
+          designed for sampled-graph (message-flow-graph) use-cases.
+          Full-graph support will be added in upcoming releases.
+
+    Parameters
+    ----------
+    in_feat : int
+        Input feature size.
+    out_feat : int
+        Output feature size.
+    num_rels : int
+        Number of relations.
+    fanout : int
+        Maximum number of sampled neighbors of an destination node;
+        i.e, maximum in degree of destination nodes
+    regularizer : str, optional
+        Which weight regularizer to use ("basis" or ``None``):
+         - "basis" is for basis-decomposition.
+         - ``None`` applies no regularization.
+        Default: ``None``.
+    num_bases : int, optional
+        Number of bases. Only used when "basis" regularizer is applied.
+        Default: ``None``.
+    bias : bool, optional
+        True if bias is added. Default: ``True``.
+    activation : callable, optional
+        Activation function. Default: ``None``.
+    self_loop : bool, optional
+        True to include self loop message. Default: ``True``.
+    dropout : float, optional
+        Dropout rate. Default: ``0.0``
+    layer_norm: bool, optional
+        True to add layer norm. Default: ``False``
+
+    Examples
+    --------
+    >>> import dgl
+    >>> from dgl.dataloading import NeighborSampler, DataLoader
+    >>> from dgl.nn import CuGraphRelGraphConv
+    >>> import torch.nn.functional as F
+    ...
+    >>> device = 'cuda'
+    >>> fanouts = [5, 6]
+    >>> sampler = NeighborSampler(fanouts)
+    >>> dataloader = DataLoader(g, train_nid, sampler,
+    ...     device=device, batch_size=1024)
+    >>> conv1 = CuGraphRelGraphConv(in_dim, h_dim, num_rels, fanouts[0],
+    ...     regularizer='basis', num_bases=10).to(device)
+    >>> conv2 = CuGraphRelGraphConv(h_dim, out_dim, num_rels, fanouts[1],
+    ...     regularizer='basis', num_bases=10).to(device)
+    >>> for input_nodes, output_nodes, blocks in dataloader:
+    ...     h = conv1(blocks[0], x, blocks[0].edata[dgl.ETYPE])
+    ...     h = F.relu(h)
+    ...     h = conv2(blocks[1], h, blocks[1].edata[dgl.ETYPE])
+    """
+    def __init__(
+        self,
+        in_feat,
+        out_feat,
+        num_rels,
+        fanout,
+        regularizer=None,
+        num_bases=None,
+        bias=True,
+        activation=None,
+        self_loop=True,
+        dropout=0.0,
+        layer_norm=False,
+    ):
+        super().__init__()
+        self.in_feat = in_feat
+        self.out_feat = out_feat
+        self.num_rels = num_rels
+        self.fanout = fanout
+
+        # regularizer (see dgl.nn.pytorch.linear.TypedLinear)
+        if regularizer is None:
+            self.W = nn.Parameter(th.Tensor(num_rels, in_feat, out_feat))
+            self.coeff = None
+        elif regularizer == "basis":
+            if num_bases is None:
+                raise ValueError(
+                    'Missing "num_bases" for basis regularization.'
+                )
+            self.W = nn.Parameter(th.Tensor(num_bases, in_feat, out_feat))
+            self.coeff = nn.Parameter(th.Tensor(num_rels, num_bases))
+            self.num_bases = num_bases
+        else:
+            raise ValueError(
+                f"Supported regularizer options: 'basis', "
+                f"but got {regularizer}"
+            )
+        self.regularizer = regularizer
+        self.reset_parameters()
+
+        # others
+        self.bias = bias
+        self.activation = activation
+        self.self_loop = self_loop
+        self.layer_norm = layer_norm
+
+        # bias
+        if self.bias:
+            self.h_bias = nn.Parameter(th.Tensor(out_feat))
+            nn.init.zeros_(self.h_bias)
+
+        # layer norm
+        if self.layer_norm:
+            self.layer_norm_weight = nn.LayerNorm(
+                out_feat, elementwise_affine=True
+            )
+
+        # weight for self_loop
+        if self.self_loop:
+            self.loop_weight = nn.Parameter(th.Tensor(in_feat, out_feat))
+            nn.init.xavier_uniform_(
+                self.loop_weight, gain=nn.init.calculate_gain("relu")
+            )
+
+        self.dropout = nn.Dropout(dropout)
+
+    def reset_parameters(self):
+        with th.no_grad():
+            if self.regularizer is None:
+                nn.init.uniform_(
+                    self.W,
+                    -1 / math.sqrt(self.in_feat),
+                    1 / math.sqrt(self.in_feat),
+                )
+            elif self.regularizer == "basis":
+                nn.init.uniform_(
+                    self.W,
+                    -1 / math.sqrt(self.in_feat),
+                    1 / math.sqrt(self.in_feat),
+                )
+                nn.init.xavier_uniform_(
+                    self.coeff, gain=nn.init.calculate_gain("relu")
+                )
+            else:
+                raise ValueError(
+                    f"Supported regularizer options: 'basis', "
+                    f"but got {self.regularizer}"
+                )
+
+    def forward(self, g, feat, etypes, norm=None):
+        r"""Forward computation.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        feat : torch.Tensor
+            A 2D tensor of node features. Shape: :math:`(|V|, D_{in})`.
+        etypes : torch.Tensor or list[int]
+            An 1D integer tensor of edge types. Shape: :math:`(|E|,)`.
+        norm : torch.Tensor, optional
+            An 1D tensor of edge norm value.  Shape: :math:`(|E|,)`.
path: python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py
text: @user1:
Ditto. The three minor issues also exist in `dgl.nn.pytorch.conv.relgraphconv`. Could you also help fix these?

diff_hunk: @@ -0,0 +1,388 @@
+"""Torch Module for Relational graph convolution layer using the
+aggregation primitives in cugraph-ops"""
+# pylint: disable=no-member, arguments-differ, invalid-name, too-many-arguments
+import math
+
+import torch as th
+from torch import nn
+
+try:
+    from pylibcugraphops.legacy.aggregators.node_level import (
+        agg_hg_basis_post_bwd_int32,
+        agg_hg_basis_post_bwd_int64,
+        agg_hg_basis_post_fwd_int32,
+        agg_hg_basis_post_fwd_int64,
+    )
+    from pylibcugraphops.legacy.structure.graph_types import (
+        message_flow_graph_hg_csr_int32,
+        message_flow_graph_hg_csr_int64,
+    )
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "dgl.nn.CuGraphRelGraphConv requires pylibcugraphops to be installed."
+    )
+
+CUDA_SM_PER_BLOCK = 49152
+
+
+class RelGraphConvAgg(th.autograd.Function):
+    r"""Custom autograd function for R-GCN aggregation layer that uses the
+    aggregation functions in cugraph-ops."""
+
+    @staticmethod
+    def forward(ctx, g, fanout, num_rels, edge_types, feat, coeff):
+        r"""Compute the forward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        fanout : int
+            Maximum in-degree of nodes.
+        num_rels : int
+            Number of relations.
+        edge_types : torch.Tensor
+            A 1D tensor of edge types. Note that pylibcugraphops only accepts
+            edge_types in int32, so any input of other integer types will be casted
+            into int32, thus introducing some overhead. Pass in int32 tensor for
+            best performance.
+        feat : torch.Tensor, dtype=torch.float32, requires_grad=True
+            Tensor of source node features. Shape: (num_src_nodes, in_feat).
+        coeff : torch.Tensor, dtype=torch.float32, requires_grad=True
+            The coefficient matrix used in basis-decomposition regularization.
+            Shape: (num_rels, num_bases). It should be set to ``None``
+            when ``regularizer=None``.
+
+        Returns
+        -------
+        agg_output : torch.Tensor, dtype=torch.float32
+            Aggregation output. Shape: (num_dst_nodes, num_rels * in_feat)
+            when ``regularizer=None``; Shape: (num_dst_nodes, num_bases * in_feat)
+            when ``regularizer='basis'``.
+        """
+        ctx.device = g.device
+        if g.idtype == th.int32:
+            mfg_csr_func = message_flow_graph_hg_csr_int32
+            agg_fwd_func = agg_hg_basis_post_fwd_int32
+        elif g.idtype == th.int64:
+            mfg_csr_func = message_flow_graph_hg_csr_int64
+            agg_fwd_func = agg_hg_basis_post_fwd_int64
+        else:
+            raise TypeError(
+                f"Supported ID type: torch.int32 or torch.int64, "
+                f"but got {g.idtype}"
+            )
+        ctx.idtype = g.idtype
+
+        _in_feat = feat.shape[-1]
+        indptr, indices, edge_ids = g.adj_sparse("csc")
+        # edge_ids is in a mixed order, need to permutate incoming etypes
+        # and stash the result for backward propagation
+        ctx.edge_types_int32 = edge_types[edge_ids.long()].int()
+        # node_types are not needed by the post-variant RGCN aggregation functions
+        _num_node_types = 0
+        _out_node_types = _in_node_types = None
+
+        mfg = mfg_csr_func(
+            fanout,
+            g.dstnodes(),
+            g.srcnodes(),
+            indptr,
+            indices,
+            _num_node_types,
+            num_rels,
+            _out_node_types,
+            _in_node_types,
+            ctx.edge_types_int32,
+        )
+        ctx.mfg = mfg
+
+        if coeff is None:
+            leading_dimension = num_rels * _in_feat
+        else:
+            _num_bases = coeff.shape[-1]
+            leading_dimension = _num_bases * _in_feat
+
+        agg_output = th.empty(
+            g.num_dst_nodes(),
+            leading_dimension,
+            dtype=th.float32,
+            device=ctx.device,
+        )
+        if coeff is None:
+            agg_fwd_func(agg_output, feat.detach(), mfg)
+        else:
+            agg_fwd_func(
+                agg_output,
+                feat.detach(),
+                mfg,
+                weights_combination=coeff.detach()
+            )
+
+        ctx.save_for_backward(feat, coeff)
+        return agg_output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        r"""Compute the backward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        grad_output : torch.Tensor, dtype=torch.float32
+            Gradient of loss function w.r.t output.
+        """
+        feat, coeff = ctx.saved_tensors
+
+        if ctx.idtype == th.int32:
+            agg_bwd_func = agg_hg_basis_post_bwd_int32
+        else:
+            agg_bwd_func = agg_hg_basis_post_bwd_int64
+
+        grad_feat = th.empty_like(feat, dtype=th.float32, device=ctx.device)
+        if coeff is None:
+            grad_coeff = None
+            agg_bwd_func(grad_feat, grad_output, feat.detach(), ctx.mfg)
+        else:
+            grad_coeff = th.empty_like(
+                coeff, dtype=th.float32, device=ctx.device
+            )
+            agg_bwd_func(
+                grad_feat,
+                grad_output,
+                feat.detach(),
+                ctx.mfg,
+                output_weight_gradient=grad_coeff,
+                weights_combination=coeff.detach()
+            )
+
+        return None, None, None, None, grad_feat, grad_coeff
+
+
+class CuGraphRelGraphConv(nn.Module):
+    r"""An accelerated relational graph convolution layer using the
+    highly-optimized aggregation primitives in cugraph-ops.
+    See :class:`dgl.nn.pytorch.conv.RelGraphConv` for mathematical model.
+
+    .. note::
+        This is an **experimental** feature.
+        Compared with :class:`dgl.nn.pytorch.conv.RelGraphConv`, this model:
+
+        * Only works on cuda devices.
+        * Only support basis-decomposition regularization.
+        * Requires an extra argument `fanout` as input, since the current model is
+          designed for sampled-graph (message-flow-graph) use-cases.
+          Full-graph support will be added in upcoming releases.
+
+    Parameters
+    ----------
+    in_feat : int
+        Input feature size.
+    out_feat : int
+        Output feature size.
+    num_rels : int
+        Number of relations.
+    fanout : int
+        Maximum number of sampled neighbors of an destination node;
+        i.e, maximum in degree of destination nodes
+    regularizer : str, optional
+        Which weight regularizer to use ("basis" or ``None``):
+         - "basis" is for basis-decomposition.
+         - ``None`` applies no regularization.
+        Default: ``None``.
+    num_bases : int, optional
+        Number of bases. Only used when "basis" regularizer is applied.
+        Default: ``None``.
+    bias : bool, optional
+        True if bias is added. Default: ``True``.
+    activation : callable, optional
+        Activation function. Default: ``None``.
+    self_loop : bool, optional
+        True to include self loop message. Default: ``True``.
+    dropout : float, optional
+        Dropout rate. Default: ``0.0``
+    layer_norm: bool, optional
+        True to add layer norm. Default: ``False``
+
+    Examples
+    --------
+    >>> import dgl
+    >>> from dgl.dataloading import NeighborSampler, DataLoader
+    >>> from dgl.nn import CuGraphRelGraphConv
+    >>> import torch.nn.functional as F
+    ...
+    >>> device = 'cuda'
+    >>> fanouts = [5, 6]
+    >>> sampler = NeighborSampler(fanouts)
+    >>> dataloader = DataLoader(g, train_nid, sampler,
+    ...     device=device, batch_size=1024)
+    >>> conv1 = CuGraphRelGraphConv(in_dim, h_dim, num_rels, fanouts[0],
+    ...     regularizer='basis', num_bases=10).to(device)
+    >>> conv2 = CuGraphRelGraphConv(h_dim, out_dim, num_rels, fanouts[1],
+    ...     regularizer='basis', num_bases=10).to(device)
+    >>> for input_nodes, output_nodes, blocks in dataloader:
+    ...     h = conv1(blocks[0], x, blocks[0].edata[dgl.ETYPE])
+    ...     h = F.relu(h)
+    ...     h = conv2(blocks[1], h, blocks[1].edata[dgl.ETYPE])
+    """
+    def __init__(
+        self,
+        in_feat,
+        out_feat,
+        num_rels,
+        fanout,
+        regularizer=None,
+        num_bases=None,
+        bias=True,
+        activation=None,
+        self_loop=True,
+        dropout=0.0,
+        layer_norm=False,
+    ):
+        super().__init__()
+        self.in_feat = in_feat
+        self.out_feat = out_feat
+        self.num_rels = num_rels
+        self.fanout = fanout
+
+        # regularizer (see dgl.nn.pytorch.linear.TypedLinear)
+        if regularizer is None:
+            self.W = nn.Parameter(th.Tensor(num_rels, in_feat, out_feat))
+            self.coeff = None
+        elif regularizer == "basis":
+            if num_bases is None:
+                raise ValueError(
+                    'Missing "num_bases" for basis regularization.'
+                )
+            self.W = nn.Parameter(th.Tensor(num_bases, in_feat, out_feat))
+            self.coeff = nn.Parameter(th.Tensor(num_rels, num_bases))
+            self.num_bases = num_bases
+        else:
+            raise ValueError(
+                f"Supported regularizer options: 'basis', "
+                f"but got {regularizer}"
+            )
+        self.regularizer = regularizer
+        self.reset_parameters()
+
+        # others
+        self.bias = bias
+        self.activation = activation
+        self.self_loop = self_loop
+        self.layer_norm = layer_norm
+
+        # bias
+        if self.bias:
+            self.h_bias = nn.Parameter(th.Tensor(out_feat))
+            nn.init.zeros_(self.h_bias)
+
+        # layer norm
+        if self.layer_norm:
+            self.layer_norm_weight = nn.LayerNorm(
+                out_feat, elementwise_affine=True
+            )
+
+        # weight for self_loop
+        if self.self_loop:
+            self.loop_weight = nn.Parameter(th.Tensor(in_feat, out_feat))
+            nn.init.xavier_uniform_(
+                self.loop_weight, gain=nn.init.calculate_gain("relu")
+            )
+
+        self.dropout = nn.Dropout(dropout)
+
+    def reset_parameters(self):
+        with th.no_grad():
+            if self.regularizer is None:
+                nn.init.uniform_(
+                    self.W,
+                    -1 / math.sqrt(self.in_feat),
+                    1 / math.sqrt(self.in_feat),
+                )
+            elif self.regularizer == "basis":
+                nn.init.uniform_(
+                    self.W,
+                    -1 / math.sqrt(self.in_feat),
+                    1 / math.sqrt(self.in_feat),
+                )
+                nn.init.xavier_uniform_(
+                    self.coeff, gain=nn.init.calculate_gain("relu")
+                )
+            else:
+                raise ValueError(
+                    f"Supported regularizer options: 'basis', "
+                    f"but got {self.regularizer}"
+                )
+
+    def forward(self, g, feat, etypes, norm=None):
+        r"""Forward computation.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        feat : torch.Tensor
+            A 2D tensor of node features. Shape: :math:`(|V|, D_{in})`.
+        etypes : torch.Tensor or list[int]
+            An 1D integer tensor of edge types. Shape: :math:`(|E|,)`.
+        norm : torch.Tensor, optional
+            An 1D tensor of edge norm value.  Shape: :math:`(|E|,)`.
+
+        Returns
+        -------
+        torch.Tensor
+            New node features. Shape: :math:`(|V|, D_{out})`.
+        """
+        _device = next(self.parameters()).device
+        if _device.type != "cuda":
+            raise RuntimeError(
+                f"dgl.nn.CuGraphRelGraphConv requires "
+                f"the model on device 'cuda', but got '{_device.type}'"
+            )
+        if _device != g.device:
+            raise RuntimeError(
+                f"Expected model and graph on the same device, "
+                f"but got '{_device}' and '{g.device}'"
+            )
+        if _device != etypes.device:
+            raise RuntimeError(
+                f"Expected model and etypes on the same device, "
+                f"but got '{_device}' and '{etypes.device}'"
+            )
+        if _device != feat.device:
+            raise RuntimeError(
+                f"Expected model and feature tensor on the same device, "
+                f"but got '{_device}' and '{feat.device}'"
+            )
path: python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py
text: @user1:
We don't do checks in L341-355 in the existing DGL NN modules. I'm not sure if they are needed.

@author:
They are actually needed. cugraph-ops consists of (almost) solely computational primitives and does not contain those checks as below in DGL,

https://github.com/dmlc/dgl/blob/e3bf18255110bffe85fccdcafcb551436e173dad/python/dgl/heterograph.py#L4121-L4124

When graphs or tensors are not properly copied to cuda devices, users will see segfault without these checks.

diff_hunk: @@ -0,0 +1,388 @@
+"""Torch Module for Relational graph convolution layer using the
+aggregation primitives in cugraph-ops"""
+# pylint: disable=no-member, arguments-differ, invalid-name, too-many-arguments
+import math
+
+import torch as th
+from torch import nn
+
+try:
+    from pylibcugraphops.legacy.aggregators.node_level import (
+        agg_hg_basis_post_bwd_int32,
+        agg_hg_basis_post_bwd_int64,
+        agg_hg_basis_post_fwd_int32,
+        agg_hg_basis_post_fwd_int64,
+    )
+    from pylibcugraphops.legacy.structure.graph_types import (
+        message_flow_graph_hg_csr_int32,
+        message_flow_graph_hg_csr_int64,
+    )
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "dgl.nn.CuGraphRelGraphConv requires pylibcugraphops to be installed."
+    )
+
+CUDA_SM_PER_BLOCK = 49152
+
+
+class RelGraphConvAgg(th.autograd.Function):
+    r"""Custom autograd function for R-GCN aggregation layer that uses the
+    aggregation functions in cugraph-ops."""
+
+    @staticmethod
+    def forward(ctx, g, fanout, num_rels, edge_types, feat, coeff):
+        r"""Compute the forward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        fanout : int
+            Maximum in-degree of nodes.
+        num_rels : int
+            Number of relations.
+        edge_types : torch.Tensor
+            A 1D tensor of edge types. Note that pylibcugraphops only accepts
+            edge_types in int32, so any input of other integer types will be casted
+            into int32, thus introducing some overhead. Pass in int32 tensor for
+            best performance.
+        feat : torch.Tensor, dtype=torch.float32, requires_grad=True
+            Tensor of source node features. Shape: (num_src_nodes, in_feat).
+        coeff : torch.Tensor, dtype=torch.float32, requires_grad=True
+            The coefficient matrix used in basis-decomposition regularization.
+            Shape: (num_rels, num_bases). It should be set to ``None``
+            when ``regularizer=None``.
+
+        Returns
+        -------
+        agg_output : torch.Tensor, dtype=torch.float32
+            Aggregation output. Shape: (num_dst_nodes, num_rels * in_feat)
+            when ``regularizer=None``; Shape: (num_dst_nodes, num_bases * in_feat)
+            when ``regularizer='basis'``.
+        """
+        ctx.device = g.device
+        if g.idtype == th.int32:
+            mfg_csr_func = message_flow_graph_hg_csr_int32
+            agg_fwd_func = agg_hg_basis_post_fwd_int32
+        elif g.idtype == th.int64:
+            mfg_csr_func = message_flow_graph_hg_csr_int64
+            agg_fwd_func = agg_hg_basis_post_fwd_int64
+        else:
+            raise TypeError(
+                f"Supported ID type: torch.int32 or torch.int64, "
+                f"but got {g.idtype}"
+            )
+        ctx.idtype = g.idtype
+
+        _in_feat = feat.shape[-1]
+        indptr, indices, edge_ids = g.adj_sparse("csc")
+        # edge_ids is in a mixed order, need to permutate incoming etypes
+        # and stash the result for backward propagation
+        ctx.edge_types_int32 = edge_types[edge_ids.long()].int()
+        # node_types are not needed by the post-variant RGCN aggregation functions
+        _num_node_types = 0
+        _out_node_types = _in_node_types = None
+
+        mfg = mfg_csr_func(
+            fanout,
+            g.dstnodes(),
+            g.srcnodes(),
+            indptr,
+            indices,
+            _num_node_types,
+            num_rels,
+            _out_node_types,
+            _in_node_types,
+            ctx.edge_types_int32,
+        )
+        ctx.mfg = mfg
+
+        if coeff is None:
+            leading_dimension = num_rels * _in_feat
+        else:
+            _num_bases = coeff.shape[-1]
+            leading_dimension = _num_bases * _in_feat
+
+        agg_output = th.empty(
+            g.num_dst_nodes(),
+            leading_dimension,
+            dtype=th.float32,
+            device=ctx.device,
+        )
+        if coeff is None:
+            agg_fwd_func(agg_output, feat.detach(), mfg)
+        else:
+            agg_fwd_func(
+                agg_output,
+                feat.detach(),
+                mfg,
+                weights_combination=coeff.detach()
+            )
+
+        ctx.save_for_backward(feat, coeff)
+        return agg_output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        r"""Compute the backward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        grad_output : torch.Tensor, dtype=torch.float32
+            Gradient of loss function w.r.t output.
+        """
+        feat, coeff = ctx.saved_tensors
+
+        if ctx.idtype == th.int32:
+            agg_bwd_func = agg_hg_basis_post_bwd_int32
+        else:
+            agg_bwd_func = agg_hg_basis_post_bwd_int64
+
+        grad_feat = th.empty_like(feat, dtype=th.float32, device=ctx.device)
+        if coeff is None:
+            grad_coeff = None
+            agg_bwd_func(grad_feat, grad_output, feat.detach(), ctx.mfg)
+        else:
+            grad_coeff = th.empty_like(
+                coeff, dtype=th.float32, device=ctx.device
+            )
+            agg_bwd_func(
+                grad_feat,
+                grad_output,
+                feat.detach(),
+                ctx.mfg,
+                output_weight_gradient=grad_coeff,
+                weights_combination=coeff.detach()
+            )
+
+        return None, None, None, None, grad_feat, grad_coeff
+
+
+class CuGraphRelGraphConv(nn.Module):
+    r"""An accelerated relational graph convolution layer using the
+    highly-optimized aggregation primitives in cugraph-ops.
+    See :class:`dgl.nn.pytorch.conv.RelGraphConv` for mathematical model.
+
+    .. note::
+        This is an **experimental** feature.
+        Compared with :class:`dgl.nn.pytorch.conv.RelGraphConv`, this model:
+
+        * Only works on cuda devices.
+        * Only support basis-decomposition regularization.
+        * Requires an extra argument `fanout` as input, since the current model is
+          designed for sampled-graph (message-flow-graph) use-cases.
+          Full-graph support will be added in upcoming releases.
+
+    Parameters
+    ----------
+    in_feat : int
+        Input feature size.
+    out_feat : int
+        Output feature size.
+    num_rels : int
+        Number of relations.
+    fanout : int
+        Maximum number of sampled neighbors of an destination node;
+        i.e, maximum in degree of destination nodes
+    regularizer : str, optional
+        Which weight regularizer to use ("basis" or ``None``):
+         - "basis" is for basis-decomposition.
+         - ``None`` applies no regularization.
+        Default: ``None``.
+    num_bases : int, optional
+        Number of bases. Only used when "basis" regularizer is applied.
+        Default: ``None``.
+    bias : bool, optional
+        True if bias is added. Default: ``True``.
+    activation : callable, optional
+        Activation function. Default: ``None``.
+    self_loop : bool, optional
+        True to include self loop message. Default: ``True``.
+    dropout : float, optional
+        Dropout rate. Default: ``0.0``
+    layer_norm: bool, optional
+        True to add layer norm. Default: ``False``
+
+    Examples
+    --------
+    >>> import dgl
+    >>> from dgl.dataloading import NeighborSampler, DataLoader
+    >>> from dgl.nn import CuGraphRelGraphConv
+    >>> import torch.nn.functional as F
+    ...
+    >>> device = 'cuda'
+    >>> fanouts = [5, 6]
+    >>> sampler = NeighborSampler(fanouts)
+    >>> dataloader = DataLoader(g, train_nid, sampler,
+    ...     device=device, batch_size=1024)
+    >>> conv1 = CuGraphRelGraphConv(in_dim, h_dim, num_rels, fanouts[0],
+    ...     regularizer='basis', num_bases=10).to(device)
+    >>> conv2 = CuGraphRelGraphConv(h_dim, out_dim, num_rels, fanouts[1],
+    ...     regularizer='basis', num_bases=10).to(device)
+    >>> for input_nodes, output_nodes, blocks in dataloader:
+    ...     h = conv1(blocks[0], x, blocks[0].edata[dgl.ETYPE])
+    ...     h = F.relu(h)
+    ...     h = conv2(blocks[1], h, blocks[1].edata[dgl.ETYPE])
+    """
+    def __init__(
+        self,
+        in_feat,
+        out_feat,
+        num_rels,
+        fanout,
+        regularizer=None,
+        num_bases=None,
+        bias=True,
+        activation=None,
+        self_loop=True,
+        dropout=0.0,
+        layer_norm=False,
+    ):
+        super().__init__()
+        self.in_feat = in_feat
+        self.out_feat = out_feat
+        self.num_rels = num_rels
+        self.fanout = fanout
+
+        # regularizer (see dgl.nn.pytorch.linear.TypedLinear)
+        if regularizer is None:
+            self.W = nn.Parameter(th.Tensor(num_rels, in_feat, out_feat))
+            self.coeff = None
+        elif regularizer == "basis":
+            if num_bases is None:
+                raise ValueError(
+                    'Missing "num_bases" for basis regularization.'
+                )
+            self.W = nn.Parameter(th.Tensor(num_bases, in_feat, out_feat))
+            self.coeff = nn.Parameter(th.Tensor(num_rels, num_bases))
+            self.num_bases = num_bases
+        else:
+            raise ValueError(
+                f"Supported regularizer options: 'basis', "
+                f"but got {regularizer}"
+            )
+        self.regularizer = regularizer
+        self.reset_parameters()
+
+        # others
+        self.bias = bias
+        self.activation = activation
+        self.self_loop = self_loop
+        self.layer_norm = layer_norm
+
+        # bias
+        if self.bias:
+            self.h_bias = nn.Parameter(th.Tensor(out_feat))
+            nn.init.zeros_(self.h_bias)
+
+        # layer norm
+        if self.layer_norm:
+            self.layer_norm_weight = nn.LayerNorm(
+                out_feat, elementwise_affine=True
+            )
+
+        # weight for self_loop
+        if self.self_loop:
+            self.loop_weight = nn.Parameter(th.Tensor(in_feat, out_feat))
+            nn.init.xavier_uniform_(
+                self.loop_weight, gain=nn.init.calculate_gain("relu")
+            )
+
+        self.dropout = nn.Dropout(dropout)
+
+    def reset_parameters(self):
+        with th.no_grad():
+            if self.regularizer is None:
+                nn.init.uniform_(
+                    self.W,
+                    -1 / math.sqrt(self.in_feat),
+                    1 / math.sqrt(self.in_feat),
+                )
+            elif self.regularizer == "basis":
+                nn.init.uniform_(
+                    self.W,
+                    -1 / math.sqrt(self.in_feat),
+                    1 / math.sqrt(self.in_feat),
+                )
+                nn.init.xavier_uniform_(
+                    self.coeff, gain=nn.init.calculate_gain("relu")
+                )
+            else:
+                raise ValueError(
+                    f"Supported regularizer options: 'basis', "
+                    f"but got {self.regularizer}"
+                )
+
+    def forward(self, g, feat, etypes, norm=None):
+        r"""Forward computation.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        feat : torch.Tensor
+            A 2D tensor of node features. Shape: :math:`(|V|, D_{in})`.
+        etypes : torch.Tensor or list[int]
+            An 1D integer tensor of edge types. Shape: :math:`(|E|,)`.
+        norm : torch.Tensor, optional
+            An 1D tensor of edge norm value.  Shape: :math:`(|E|,)`.
+
+        Returns
+        -------
+        torch.Tensor
+            New node features. Shape: :math:`(|V|, D_{out})`.
+        """
+        _device = next(self.parameters()).device
+        if _device.type != "cuda":
+            raise RuntimeError(
+                f"dgl.nn.CuGraphRelGraphConv requires "
+                f"the model on device 'cuda', but got '{_device.type}'"
path: python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py
text: @user1:
```suggestion
                f"the model to be on device 'cuda', but got '{_device.type}'"
```

diff_hunk: @@ -0,0 +1,388 @@
+"""Torch Module for Relational graph convolution layer using the
+aggregation primitives in cugraph-ops"""
+# pylint: disable=no-member, arguments-differ, invalid-name, too-many-arguments
+import math
+
+import torch as th
+from torch import nn
+
+try:
+    from pylibcugraphops.legacy.aggregators.node_level import (
+        agg_hg_basis_post_bwd_int32,
+        agg_hg_basis_post_bwd_int64,
+        agg_hg_basis_post_fwd_int32,
+        agg_hg_basis_post_fwd_int64,
+    )
+    from pylibcugraphops.legacy.structure.graph_types import (
+        message_flow_graph_hg_csr_int32,
+        message_flow_graph_hg_csr_int64,
+    )
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "dgl.nn.CuGraphRelGraphConv requires pylibcugraphops to be installed."
+    )
+
+CUDA_SM_PER_BLOCK = 49152
+
+
+class RelGraphConvAgg(th.autograd.Function):
+    r"""Custom autograd function for R-GCN aggregation layer that uses the
+    aggregation functions in cugraph-ops."""
+
+    @staticmethod
+    def forward(ctx, g, fanout, num_rels, edge_types, feat, coeff):
+        r"""Compute the forward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        fanout : int
+            Maximum in-degree of nodes.
+        num_rels : int
+            Number of relations.
+        edge_types : torch.Tensor
+            A 1D tensor of edge types. Note that pylibcugraphops only accepts
+            edge_types in int32, so any input of other integer types will be casted
+            into int32, thus introducing some overhead. Pass in int32 tensor for
+            best performance.
+        feat : torch.Tensor, dtype=torch.float32, requires_grad=True
+            Tensor of source node features. Shape: (num_src_nodes, in_feat).
+        coeff : torch.Tensor, dtype=torch.float32, requires_grad=True
+            The coefficient matrix used in basis-decomposition regularization.
+            Shape: (num_rels, num_bases). It should be set to ``None``
+            when ``regularizer=None``.
+
+        Returns
+        -------
+        agg_output : torch.Tensor, dtype=torch.float32
+            Aggregation output. Shape: (num_dst_nodes, num_rels * in_feat)
+            when ``regularizer=None``; Shape: (num_dst_nodes, num_bases * in_feat)
+            when ``regularizer='basis'``.
+        """
+        ctx.device = g.device
+        if g.idtype == th.int32:
+            mfg_csr_func = message_flow_graph_hg_csr_int32
+            agg_fwd_func = agg_hg_basis_post_fwd_int32
+        elif g.idtype == th.int64:
+            mfg_csr_func = message_flow_graph_hg_csr_int64
+            agg_fwd_func = agg_hg_basis_post_fwd_int64
+        else:
+            raise TypeError(
+                f"Supported ID type: torch.int32 or torch.int64, "
+                f"but got {g.idtype}"
+            )
+        ctx.idtype = g.idtype
+
+        _in_feat = feat.shape[-1]
+        indptr, indices, edge_ids = g.adj_sparse("csc")
+        # edge_ids is in a mixed order, need to permutate incoming etypes
+        # and stash the result for backward propagation
+        ctx.edge_types_int32 = edge_types[edge_ids.long()].int()
+        # node_types are not needed by the post-variant RGCN aggregation functions
+        _num_node_types = 0
+        _out_node_types = _in_node_types = None
+
+        mfg = mfg_csr_func(
+            fanout,
+            g.dstnodes(),
+            g.srcnodes(),
+            indptr,
+            indices,
+            _num_node_types,
+            num_rels,
+            _out_node_types,
+            _in_node_types,
+            ctx.edge_types_int32,
+        )
+        ctx.mfg = mfg
+
+        if coeff is None:
+            leading_dimension = num_rels * _in_feat
+        else:
+            _num_bases = coeff.shape[-1]
+            leading_dimension = _num_bases * _in_feat
+
+        agg_output = th.empty(
+            g.num_dst_nodes(),
+            leading_dimension,
+            dtype=th.float32,
+            device=ctx.device,
+        )
+        if coeff is None:
+            agg_fwd_func(agg_output, feat.detach(), mfg)
+        else:
+            agg_fwd_func(
+                agg_output,
+                feat.detach(),
+                mfg,
+                weights_combination=coeff.detach()
+            )
+
+        ctx.save_for_backward(feat, coeff)
+        return agg_output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        r"""Compute the backward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        grad_output : torch.Tensor, dtype=torch.float32
+            Gradient of loss function w.r.t output.
+        """
+        feat, coeff = ctx.saved_tensors
+
+        if ctx.idtype == th.int32:
+            agg_bwd_func = agg_hg_basis_post_bwd_int32
+        else:
+            agg_bwd_func = agg_hg_basis_post_bwd_int64
+
+        grad_feat = th.empty_like(feat, dtype=th.float32, device=ctx.device)
+        if coeff is None:
+            grad_coeff = None
+            agg_bwd_func(grad_feat, grad_output, feat.detach(), ctx.mfg)
+        else:
+            grad_coeff = th.empty_like(
+                coeff, dtype=th.float32, device=ctx.device
+            )
+            agg_bwd_func(
+                grad_feat,
+                grad_output,
+                feat.detach(),
+                ctx.mfg,
+                output_weight_gradient=grad_coeff,
+                weights_combination=coeff.detach()
+            )
+
+        return None, None, None, None, grad_feat, grad_coeff
+
+
+class CuGraphRelGraphConv(nn.Module):
+    r"""An accelerated relational graph convolution layer using the
+    highly-optimized aggregation primitives in cugraph-ops.
+    See :class:`dgl.nn.pytorch.conv.RelGraphConv` for mathematical model.
+
+    .. note::
+        This is an **experimental** feature.
+        Compared with :class:`dgl.nn.pytorch.conv.RelGraphConv`, this model:
+
+        * Only works on cuda devices.
+        * Only support basis-decomposition regularization.
+        * Requires an extra argument `fanout` as input, since the current model is
+          designed for sampled-graph (message-flow-graph) use-cases.
+          Full-graph support will be added in upcoming releases.
+
+    Parameters
+    ----------
+    in_feat : int
+        Input feature size.
+    out_feat : int
+        Output feature size.
+    num_rels : int
+        Number of relations.
+    fanout : int
+        Maximum number of sampled neighbors of an destination node;
+        i.e, maximum in degree of destination nodes
+    regularizer : str, optional
+        Which weight regularizer to use ("basis" or ``None``):
+         - "basis" is for basis-decomposition.
+         - ``None`` applies no regularization.
+        Default: ``None``.
+    num_bases : int, optional
+        Number of bases. Only used when "basis" regularizer is applied.
+        Default: ``None``.
+    bias : bool, optional
+        True if bias is added. Default: ``True``.
+    activation : callable, optional
+        Activation function. Default: ``None``.
+    self_loop : bool, optional
+        True to include self loop message. Default: ``True``.
+    dropout : float, optional
+        Dropout rate. Default: ``0.0``
+    layer_norm: bool, optional
+        True to add layer norm. Default: ``False``
+
+    Examples
+    --------
+    >>> import dgl
+    >>> from dgl.dataloading import NeighborSampler, DataLoader
+    >>> from dgl.nn import CuGraphRelGraphConv
+    >>> import torch.nn.functional as F
+    ...
+    >>> device = 'cuda'
+    >>> fanouts = [5, 6]
+    >>> sampler = NeighborSampler(fanouts)
+    >>> dataloader = DataLoader(g, train_nid, sampler,
+    ...     device=device, batch_size=1024)
+    >>> conv1 = CuGraphRelGraphConv(in_dim, h_dim, num_rels, fanouts[0],
+    ...     regularizer='basis', num_bases=10).to(device)
+    >>> conv2 = CuGraphRelGraphConv(h_dim, out_dim, num_rels, fanouts[1],
+    ...     regularizer='basis', num_bases=10).to(device)
+    >>> for input_nodes, output_nodes, blocks in dataloader:
+    ...     h = conv1(blocks[0], x, blocks[0].edata[dgl.ETYPE])
+    ...     h = F.relu(h)
+    ...     h = conv2(blocks[1], h, blocks[1].edata[dgl.ETYPE])
+    """
+    def __init__(
+        self,
+        in_feat,
+        out_feat,
+        num_rels,
+        fanout,
+        regularizer=None,
+        num_bases=None,
+        bias=True,
+        activation=None,
+        self_loop=True,
+        dropout=0.0,
+        layer_norm=False,
+    ):
+        super().__init__()
+        self.in_feat = in_feat
+        self.out_feat = out_feat
+        self.num_rels = num_rels
+        self.fanout = fanout
+
+        # regularizer (see dgl.nn.pytorch.linear.TypedLinear)
+        if regularizer is None:
+            self.W = nn.Parameter(th.Tensor(num_rels, in_feat, out_feat))
+            self.coeff = None
+        elif regularizer == "basis":
+            if num_bases is None:
+                raise ValueError(
+                    'Missing "num_bases" for basis regularization.'
+                )
+            self.W = nn.Parameter(th.Tensor(num_bases, in_feat, out_feat))
+            self.coeff = nn.Parameter(th.Tensor(num_rels, num_bases))
+            self.num_bases = num_bases
+        else:
+            raise ValueError(
+                f"Supported regularizer options: 'basis', "
+                f"but got {regularizer}"
+            )
+        self.regularizer = regularizer
+        self.reset_parameters()
+
+        # others
+        self.bias = bias
+        self.activation = activation
+        self.self_loop = self_loop
+        self.layer_norm = layer_norm
+
+        # bias
+        if self.bias:
+            self.h_bias = nn.Parameter(th.Tensor(out_feat))
+            nn.init.zeros_(self.h_bias)
+
+        # layer norm
+        if self.layer_norm:
+            self.layer_norm_weight = nn.LayerNorm(
+                out_feat, elementwise_affine=True
+            )
+
+        # weight for self_loop
+        if self.self_loop:
+            self.loop_weight = nn.Parameter(th.Tensor(in_feat, out_feat))
+            nn.init.xavier_uniform_(
+                self.loop_weight, gain=nn.init.calculate_gain("relu")
+            )
+
+        self.dropout = nn.Dropout(dropout)
+
+    def reset_parameters(self):
+        with th.no_grad():
+            if self.regularizer is None:
+                nn.init.uniform_(
+                    self.W,
+                    -1 / math.sqrt(self.in_feat),
+                    1 / math.sqrt(self.in_feat),
+                )
+            elif self.regularizer == "basis":
+                nn.init.uniform_(
+                    self.W,
+                    -1 / math.sqrt(self.in_feat),
+                    1 / math.sqrt(self.in_feat),
+                )
+                nn.init.xavier_uniform_(
+                    self.coeff, gain=nn.init.calculate_gain("relu")
+                )
+            else:
+                raise ValueError(
+                    f"Supported regularizer options: 'basis', "
+                    f"but got {self.regularizer}"
+                )
+
+    def forward(self, g, feat, etypes, norm=None):
+        r"""Forward computation.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        feat : torch.Tensor
+            A 2D tensor of node features. Shape: :math:`(|V|, D_{in})`.
+        etypes : torch.Tensor or list[int]
+            An 1D integer tensor of edge types. Shape: :math:`(|E|,)`.
path: python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py
text: @user1:
```suggestion
            A 1D integer tensor of edge types. Shape: :math:`(|E|,)`.
```

diff_hunk: @@ -0,0 +1,388 @@
+"""Torch Module for Relational graph convolution layer using the
+aggregation primitives in cugraph-ops"""
+# pylint: disable=no-member, arguments-differ, invalid-name, too-many-arguments
+import math
+
+import torch as th
+from torch import nn
+
+try:
+    from pylibcugraphops.legacy.aggregators.node_level import (
+        agg_hg_basis_post_bwd_int32,
+        agg_hg_basis_post_bwd_int64,
+        agg_hg_basis_post_fwd_int32,
+        agg_hg_basis_post_fwd_int64,
+    )
+    from pylibcugraphops.legacy.structure.graph_types import (
+        message_flow_graph_hg_csr_int32,
+        message_flow_graph_hg_csr_int64,
+    )
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "dgl.nn.CuGraphRelGraphConv requires pylibcugraphops to be installed."
+    )
+
+CUDA_SM_PER_BLOCK = 49152
+
+
+class RelGraphConvAgg(th.autograd.Function):
+    r"""Custom autograd function for R-GCN aggregation layer that uses the
+    aggregation functions in cugraph-ops."""
+
+    @staticmethod
+    def forward(ctx, g, fanout, num_rels, edge_types, feat, coeff):
+        r"""Compute the forward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        fanout : int
+            Maximum in-degree of nodes.
+        num_rels : int
+            Number of relations.
+        edge_types : torch.Tensor
+            A 1D tensor of edge types. Note that pylibcugraphops only accepts
+            edge_types in int32, so any input of other integer types will be casted
+            into int32, thus introducing some overhead. Pass in int32 tensor for
+            best performance.
+        feat : torch.Tensor, dtype=torch.float32, requires_grad=True
+            Tensor of source node features. Shape: (num_src_nodes, in_feat).
path: python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py
text: @user1:
Again, there can be inconsistencies between parameter description in this class and that in `CuGraphRelGraphConv`.

diff_hunk: @@ -0,0 +1,388 @@
+"""Torch Module for Relational graph convolution layer using the
+aggregation primitives in cugraph-ops"""
+# pylint: disable=no-member, arguments-differ, invalid-name, too-many-arguments
+import math
+
+import torch as th
+from torch import nn
+
+try:
+    from pylibcugraphops.legacy.aggregators.node_level import (
+        agg_hg_basis_post_bwd_int32,
+        agg_hg_basis_post_bwd_int64,
+        agg_hg_basis_post_fwd_int32,
+        agg_hg_basis_post_fwd_int64,
+    )
+    from pylibcugraphops.legacy.structure.graph_types import (
+        message_flow_graph_hg_csr_int32,
+        message_flow_graph_hg_csr_int64,
+    )
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "dgl.nn.CuGraphRelGraphConv requires pylibcugraphops to be installed."
+    )
+
+CUDA_SM_PER_BLOCK = 49152
+
+
+class RelGraphConvAgg(th.autograd.Function):
+    r"""Custom autograd function for R-GCN aggregation layer that uses the
+    aggregation functions in cugraph-ops."""
+
+    @staticmethod
+    def forward(ctx, g, fanout, num_rels, edge_types, feat, coeff):
+        r"""Compute the forward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        fanout : int
+            Maximum in-degree of nodes.
+        num_rels : int
+            Number of relations.
+        edge_types : torch.Tensor
+            A 1D tensor of edge types. Note that pylibcugraphops only accepts
+            edge_types in int32, so any input of other integer types will be casted
+            into int32, thus introducing some overhead. Pass in int32 tensor for
+            best performance.
+        feat : torch.Tensor, dtype=torch.float32, requires_grad=True
+            Tensor of source node features. Shape: (num_src_nodes, in_feat).
+        coeff : torch.Tensor, dtype=torch.float32, requires_grad=True
path: python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py
text: @user1:
You can omit `dtype=torch.float32, requires_grad=True` since this is created in `CuGraphRelGraphConv`.

@author:
Fixed.

diff_hunk: @@ -0,0 +1,388 @@
+"""Torch Module for Relational graph convolution layer using the
+aggregation primitives in cugraph-ops"""
+# pylint: disable=no-member, arguments-differ, invalid-name, too-many-arguments
+import math
+
+import torch as th
+from torch import nn
+
+try:
+    from pylibcugraphops.legacy.aggregators.node_level import (
+        agg_hg_basis_post_bwd_int32,
+        agg_hg_basis_post_bwd_int64,
+        agg_hg_basis_post_fwd_int32,
+        agg_hg_basis_post_fwd_int64,
+    )
+    from pylibcugraphops.legacy.structure.graph_types import (
+        message_flow_graph_hg_csr_int32,
+        message_flow_graph_hg_csr_int64,
+    )
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "dgl.nn.CuGraphRelGraphConv requires pylibcugraphops to be installed."
+    )
+
+CUDA_SM_PER_BLOCK = 49152
+
+
+class RelGraphConvAgg(th.autograd.Function):
+    r"""Custom autograd function for R-GCN aggregation layer that uses the
+    aggregation functions in cugraph-ops."""
+
+    @staticmethod
+    def forward(ctx, g, fanout, num_rels, edge_types, feat, coeff):
+        r"""Compute the forward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        fanout : int
+            Maximum in-degree of nodes.
+        num_rels : int
+            Number of relations.
+        edge_types : torch.Tensor
+            A 1D tensor of edge types. Note that pylibcugraphops only accepts
+            edge_types in int32, so any input of other integer types will be casted
+            into int32, thus introducing some overhead. Pass in int32 tensor for
+            best performance.
+        feat : torch.Tensor, dtype=torch.float32, requires_grad=True
+            Tensor of source node features. Shape: (num_src_nodes, in_feat).
+        coeff : torch.Tensor, dtype=torch.float32, requires_grad=True
+            The coefficient matrix used in basis-decomposition regularization.
+            Shape: (num_rels, num_bases). It should be set to ``None``
+            when ``regularizer=None``.
+
+        Returns
+        -------
+        agg_output : torch.Tensor, dtype=torch.float32
+            Aggregation output. Shape: (num_dst_nodes, num_rels * in_feat)
+            when ``regularizer=None``; Shape: (num_dst_nodes, num_bases * in_feat)
+            when ``regularizer='basis'``.
+        """
+        ctx.device = g.device
path: python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py
text: @user1:
This seems redundant since you already assume that in `CuGraphRelGraphConv`.

@author:
Good catch. Fixed.

diff_hunk: @@ -0,0 +1,388 @@
+"""Torch Module for Relational graph convolution layer using the
+aggregation primitives in cugraph-ops"""
+# pylint: disable=no-member, arguments-differ, invalid-name, too-many-arguments
+import math
+
+import torch as th
+from torch import nn
+
+try:
+    from pylibcugraphops.legacy.aggregators.node_level import (
+        agg_hg_basis_post_bwd_int32,
+        agg_hg_basis_post_bwd_int64,
+        agg_hg_basis_post_fwd_int32,
+        agg_hg_basis_post_fwd_int64,
+    )
+    from pylibcugraphops.legacy.structure.graph_types import (
+        message_flow_graph_hg_csr_int32,
+        message_flow_graph_hg_csr_int64,
+    )
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "dgl.nn.CuGraphRelGraphConv requires pylibcugraphops to be installed."
+    )
+
+CUDA_SM_PER_BLOCK = 49152
+
+
+class RelGraphConvAgg(th.autograd.Function):
+    r"""Custom autograd function for R-GCN aggregation layer that uses the
+    aggregation functions in cugraph-ops."""
+
+    @staticmethod
+    def forward(ctx, g, fanout, num_rels, edge_types, feat, coeff):
+        r"""Compute the forward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        fanout : int
+            Maximum in-degree of nodes.
+        num_rels : int
+            Number of relations.
+        edge_types : torch.Tensor
+            A 1D tensor of edge types. Note that pylibcugraphops only accepts
+            edge_types in int32, so any input of other integer types will be casted
+            into int32, thus introducing some overhead. Pass in int32 tensor for
+            best performance.
+        feat : torch.Tensor, dtype=torch.float32, requires_grad=True
+            Tensor of source node features. Shape: (num_src_nodes, in_feat).
+        coeff : torch.Tensor, dtype=torch.float32, requires_grad=True
+            The coefficient matrix used in basis-decomposition regularization.
+            Shape: (num_rels, num_bases). It should be set to ``None``
+            when ``regularizer=None``.
+
+        Returns
+        -------
+        agg_output : torch.Tensor, dtype=torch.float32
path: python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py
text: @user1:
No need to specify `dtype`

diff_hunk: @@ -0,0 +1,388 @@
+"""Torch Module for Relational graph convolution layer using the
+aggregation primitives in cugraph-ops"""
+# pylint: disable=no-member, arguments-differ, invalid-name, too-many-arguments
+import math
+
+import torch as th
+from torch import nn
+
+try:
+    from pylibcugraphops.legacy.aggregators.node_level import (
+        agg_hg_basis_post_bwd_int32,
+        agg_hg_basis_post_bwd_int64,
+        agg_hg_basis_post_fwd_int32,
+        agg_hg_basis_post_fwd_int64,
+    )
+    from pylibcugraphops.legacy.structure.graph_types import (
+        message_flow_graph_hg_csr_int32,
+        message_flow_graph_hg_csr_int64,
+    )
+except ModuleNotFoundError:
+    raise ModuleNotFoundError(
+        "dgl.nn.CuGraphRelGraphConv requires pylibcugraphops to be installed."
+    )
+
+CUDA_SM_PER_BLOCK = 49152
+
+
+class RelGraphConvAgg(th.autograd.Function):
+    r"""Custom autograd function for R-GCN aggregation layer that uses the
+    aggregation functions in cugraph-ops."""
+
+    @staticmethod
+    def forward(ctx, g, fanout, num_rels, edge_types, feat, coeff):
+        r"""Compute the forward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        g : DGLGraph
+            The graph.
+        fanout : int
+            Maximum in-degree of nodes.
+        num_rels : int
+            Number of relations.
+        edge_types : torch.Tensor
+            A 1D tensor of edge types. Note that pylibcugraphops only accepts
+            edge_types in int32, so any input of other integer types will be casted
+            into int32, thus introducing some overhead. Pass in int32 tensor for
+            best performance.
+        feat : torch.Tensor, dtype=torch.float32, requires_grad=True
+            Tensor of source node features. Shape: (num_src_nodes, in_feat).
+        coeff : torch.Tensor, dtype=torch.float32, requires_grad=True
+            The coefficient matrix used in basis-decomposition regularization.
+            Shape: (num_rels, num_bases). It should be set to ``None``
+            when ``regularizer=None``.
+
+        Returns
+        -------
+        agg_output : torch.Tensor, dtype=torch.float32
+            Aggregation output. Shape: (num_dst_nodes, num_rels * in_feat)
+            when ``regularizer=None``; Shape: (num_dst_nodes, num_bases * in_feat)
+            when ``regularizer='basis'``.
+        """
+        ctx.device = g.device
+        if g.idtype == th.int32:
+            mfg_csr_func = message_flow_graph_hg_csr_int32
+            agg_fwd_func = agg_hg_basis_post_fwd_int32
+        elif g.idtype == th.int64:
+            mfg_csr_func = message_flow_graph_hg_csr_int64
+            agg_fwd_func = agg_hg_basis_post_fwd_int64
+        else:
+            raise TypeError(
+                f"Supported ID type: torch.int32 or torch.int64, "
+                f"but got {g.idtype}"
+            )
+        ctx.idtype = g.idtype
+
+        _in_feat = feat.shape[-1]
+        indptr, indices, edge_ids = g.adj_sparse("csc")
+        # edge_ids is in a mixed order, need to permutate incoming etypes
+        # and stash the result for backward propagation
+        ctx.edge_types_int32 = edge_types[edge_ids.long()].int()
+        # node_types are not needed by the post-variant RGCN aggregation functions
+        _num_node_types = 0
+        _out_node_types = _in_node_types = None
+
+        mfg = mfg_csr_func(
+            fanout,
+            g.dstnodes(),
+            g.srcnodes(),
+            indptr,
+            indices,
+            _num_node_types,
+            num_rels,
+            _out_node_types,
+            _in_node_types,
+            ctx.edge_types_int32,
+        )
+        ctx.mfg = mfg
+
+        if coeff is None:
+            leading_dimension = num_rels * _in_feat
+        else:
+            _num_bases = coeff.shape[-1]
+            leading_dimension = _num_bases * _in_feat
+
+        agg_output = th.empty(
+            g.num_dst_nodes(),
+            leading_dimension,
+            dtype=th.float32,
+            device=ctx.device,
+        )
+        if coeff is None:
+            agg_fwd_func(agg_output, feat.detach(), mfg)
+        else:
+            agg_fwd_func(
+                agg_output,
+                feat.detach(),
+                mfg,
+                weights_combination=coeff.detach()
+            )
+
+        ctx.save_for_backward(feat, coeff)
+        return agg_output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        r"""Compute the backward pass of R-GCN aggregation layer.
+
+        Parameters
+        ----------
+        grad_output : torch.Tensor, dtype=torch.float32
+            Gradient of loss function w.r.t output.
path: python/dgl/nn/pytorch/conv/cugraph_relgraphconv.py
text: @user1:
Ditto for `ctx` and removing dtype of `grad_output`.

@author:
dtype removed. Regarding ctx, see comments above.