test conda build

build.sh

@@ -31,6 +31,7 @@ VALIDARGS="
    cugraph-service
    cugraph-pyg
    cugraph-dgl
+   cugraph-equivariant
    nx-cugraph
    cpp-mgtests
    cpp-mtmgtests
@@ -60,6 +61,7 @@ HELP="$0 [<target> ...] [<flag> ...]
    cugraph-service            - build the cugraph-service_client and cugraph-service_server Python package
    cugraph-pyg                - build the cugraph-pyg Python package
    cugraph-dgl                - build the cugraph-dgl extensions for DGL
+   cugraph-equivariant        - build the cugraph-equivariant Python package
    nx-cugraph                 - build the nx-cugraph Python package
    cpp-mgtests                - build libcugraph and libcugraph_etl MG tests. Builds MPI communicator, adding MPI as a dependency.
    cpp-mtmgtests              - build libcugraph MTMG tests. Adds UCX as a dependency (temporary).
@@ -219,7 +221,7 @@ if hasArg uninstall; then
     # removes the latest one and leaves the others installed. build.sh uninstall
     # can be run multiple times to remove all of them, but that is not obvious.
     pip uninstall -y pylibcugraph cugraph cugraph-service-client cugraph-service-server \
-        cugraph-dgl cugraph-pyg nx-cugraph
+        cugraph-dgl cugraph-pyg cugraph-equivariant nx-cugraph
 fi
 
 if hasArg clean; then
@@ -394,6 +396,15 @@ if hasArg cugraph-dgl || hasArg all; then
     fi
 fi
 
+# Build and install the cugraph-equivariant Python package
+if hasArg cugraph-equivariant || hasArg all; then
+    if hasArg --clean; then
+        cleanPythonDir ${REPODIR}/python/cugraph-equivariant
+    else
+        python ${PYTHON_ARGS_FOR_INSTALL} ${REPODIR}/python/cugraph-equivariant
+    fi
+fi
+
 # Build and install the nx-cugraph Python package
 if hasArg nx-cugraph || hasArg all; then
     if hasArg --clean; then

ci/build_python.sh

@@ -19,7 +19,7 @@ echo "${version}" > VERSION
 rapids-logger "Begin py build"
 
 package_dir="python"
-for package_name in pylibcugraph cugraph nx-cugraph cugraph-pyg cugraph-dgl; do 
+for package_name in pylibcugraph cugraph nx-cugraph cugraph-pyg cugraph-dgl; do
   underscore_package_name=$(echo "${package_name}" | tr "-" "_")
   sed -i "/^__git_commit__/ s/= .*/= \"${git_commit}\"/g" "${package_dir}/${package_name}/${underscore_package_name}/_version.py"
 done
@@ -85,4 +85,10 @@ if [[ ${RAPIDS_CUDA_MAJOR} == "11" ]]; then
     conda/recipes/cugraph-dgl
 fi
 
+rapids-conda-retry mambabuild \
+  --no-test \
+  --channel "${CPP_CHANNEL}" \
+  --channel "${RAPIDS_CONDA_BLD_OUTPUT_DIR}" \
+  conda/recipes/cugraph-equivariant
+
 rapids-upload-conda-to-s3 python

conda/recipes/cugraph-equivariant/build.sh

@@ -0,0 +1,7 @@
+#!/usr/bin/env bash
+
+# Copyright (c) 2023, NVIDIA CORPORATION.
+
+# This assumes the script is executed from the root of the repo directory
+
+./build.sh cugraph-equivariant

conda/recipes/cugraph-equivariant/meta.yaml

@@ -0,0 +1,39 @@
+# Copyright (c) 2023, NVIDIA CORPORATION.
+
+{% set version = environ['RAPIDS_PACKAGE_VERSION'].lstrip('v') + environ.get('VERSION_SUFFIX', '') %}
+{% set minor_version = version.split('.')[0] + '.' + version.split('.')[1] %}
+{% set py_version = environ['CONDA_PY'] %}
+{% set date_string = environ['RAPIDS_DATE_STRING'] %}
+
+package:
+  name: cugraph-equivariant
+  version: {{ version }}
+
+source:
+  path: ../../..
+
+build:
+  number: {{ GIT_DESCRIBE_NUMBER }}
+  build:
+      number: {{ GIT_DESCRIBE_NUMBER }}
+      string: py{{ py_version }}_{{ date_string }}_{{ GIT_DESCRIBE_HASH }}_{{ GIT_DESCRIBE_NUMBER }}
+
+requirements:
+  host:
+    - python
+  run:
+    - e3nn >=0.5.1
+    - pylibcugraphops ={{ minor_version }}
+    - pytorch >=2.0
+    - python
+
+tests:
+  imports:
+    - cugraph_equivariant
+
+about:
+  home: https://rapids.ai/
+  dev_url: https://github.com/rapidsai/cugraph
+  license: Apache-2.0
+  license_file: ../../../LICENSE
+  summary: GPU-accelerated equivariant convolutional layers.

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

@@ -14,17 +14,18 @@
 import pytest
 import torch
 
+
 @pytest.fixture
 def example_scatter_data():
     src_feat = torch.Tensor([3, 1, 0, 1, 1, 2])
     dst_indices = torch.Tensor([0, 1, 2, 2, 3, 1])
 
     results = {
-        "sum": torch.Tensor([3., 3., 1., 1.]),
-        "mean": torch.Tensor([3., 1.5, 0.5, 1.]),
-        "prod": torch.Tensor([3., 2., 0., 1.]),
-        "amax": torch.Tensor([3., 2., 1., 1.]),
-        "amin": torch.Tensor([3., 1., 0., 1.]),
+        "sum": torch.Tensor([3.0, 3.0, 1.0, 1.0]),
+        "mean": torch.Tensor([3.0, 1.5, 0.5, 1.0]),
+        "prod": torch.Tensor([3.0, 2.0, 0.0, 1.0]),
+        "amax": torch.Tensor([3.0, 2.0, 1.0, 1.0]),
+        "amin": torch.Tensor([3.0, 1.0, 0.0, 1.0]),
     }
 
     return src_feat, dst_indices, results

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

@@ -15,6 +15,7 @@
 import torch
 from cugraph_equivariant.utils import scatter_reduce
 
+
 @pytest.mark.parametrize("reduce", ["sum", "mean", "prod", "amax", "amin"])
 def test_scatter_reduce(example_scatter_data, reduce):
     device = torch.device("cuda:0")