### VSCode Debug Launch Configuration

Source: https://github.com/rapidsai/cugraph-gnn/blob/main/readme_pages/CONTRIBUTING.md

An example of a VSCode launch configuration for debugging cuGraph. This configuration, typically found in a '.vscode/launch.json' file, enables debugging directly within the VSCode IDE.

```json
{
  "version": "0.2.0",
  "configurations": [
    {
      "name": "Python: Current File",
      "type": "python",
      "request": "launch",
      "program": "${file}",
      "console": "integratedTerminal",
      "justMyCode": false
    }
  ]
}
```

--------------------------------

### Project Setup and Language Configuration

Source: https://github.com/rapidsai/cugraph-gnn/blob/main/python/libwholegraph/CMakeLists.txt

Configures the project named LIBWHOLEGRAPH with a specified version and enables CXX and CUDA languages. This is a standard CMake project setup command.

```cmake
project(
  LIBWHOLEGRAPH
  VERSION "${RAPIDS_VERSION}"
  LANGUAGES CXX CUDA
)
```

--------------------------------

### Define Project Namespaces

Source: https://github.com/rapidsai/cugraph-gnn/blob/main/cpp/docs/DEVELOPER_GUIDE.md

Illustrates the standard namespace organization for wholegraph, including public APIs, internal detail namespaces for cross-file utilities, and anonymous namespaces for file-local scope.

```C++
namespace wholegraph{
   void public_function(...);
} // namespace wholegraph

// some_utilities.hpp
namespace wholegraph{
namespace detail{
void reusable_helper_function(...);
} // namespace detail
} // namespace wholegraph

// some_file.cpp
namespace{
void isolated_helper_function(...);
} // anonymous namespace
```

--------------------------------

### Cython Core Initialization for cugraph-gnn

Source: https://github.com/rapidsai/cugraph-gnn/blob/main/python/pylibwholegraph/CMakeLists.txt

Initializes the Cython core for the cugraph-gnn project. This setup is necessary for building Python extensions using Cython, enabling efficient integration of C++/CUDA code with Python.

```cmake
include(rapids-cython-core)
rapids_cython_init()
```

--------------------------------

### Create LinkNeighborLoader for Link Prediction in PyTorch

Source: https://context7.com/rapidsai/cugraph-gnn/llms.txt

Shows how to set up and utilize the LinkNeighborLoader for link prediction tasks. This loader samples neighbors starting from edges and includes negative sampling. It requires setting up stores, defining edge label indices, and configuring the negative sampling strategy.

```python
import torch
from cugraph_pyg.data import GraphStore, FeatureStore
from cugraph_pyg.loader import LinkNeighborLoader
from torch_geometric.sampler import NegativeSampling

# Setup stores
graph_store = GraphStore()
feature_store = FeatureStore()
data = (feature_store, graph_store)

# Define edge label index for training (positive edges)
edge_label_index = torch.stack([
    torch.randint(0, 10000, (5000,)),  # source nodes
    torch.randint(0, 10000, (5000,)),  # target nodes
], dim=0)

# Create LinkNeighborLoader with negative sampling
loader = LinkNeighborLoader(
    data,
    num_neighbors=[10, 5],          # fanout per hop
    edge_label_index=edge_label_index,
    edge_label=torch.ones(5000),    # all positive edges
    batch_size=512,
    shuffle=True,
    neg_sampling=NegativeSampling(
        mode="binary",               # binary or triplet
        amount=1.0,                  # ratio of negative to positive samples
    ),
)

# Training loop for link prediction (assuming 'model', 'F', and 'optimizer' are defined)
# for batch in loader:
#     batch = batch.to("cuda")

#     # batch.edge_label contains 1s for positive, 0s for negative edges
#     # batch.edge_label_index contains the edges being predicted
#     z = model.encode(batch.x, batch.edge_index)

#     pred = model.decode(z, batch.edge_label_index)
#     loss = F.binary_cross_entropy_with_logits(pred, batch.edge_label)

#     loss.backward()
#     optimizer.step()
```

--------------------------------

### WholeGraph Initialization and Communication Setup (pylibwholegraph)

Source: https://context7.com/rapidsai/cugraph-gnn/llms.txt

Initializes the WholeGraph distributed memory system and cuGraph communications for multi-GPU training. This involves setting up PyTorch distributed, CUDA devices, and communication backends.

```python
import torch
import torch.distributed as dist
import os
from pylibwholegraph.torch.initialize import init as wm_init, finalize as wm_finalize
from pylibcugraph.comms import cugraph_comms_init, cugraph_comms_create_unique_id

# Initialize PyTorch distributed
dist.init_process_group("nccl")
world_size = dist.get_world_size()
global_rank = dist.get_rank()
local_rank = int(os.environ["LOCAL_RANK"])

# Set CUDA device
torch.cuda.set_device(local_rank)

# Create cuGraph unique ID (only rank 0)
if global_rank == 0:
    cugraph_id = [cugraph_comms_create_unique_id()]
else:
    cugraph_id = [None]
dist.broadcast_object_list(cugraph_id, src=0)

# Initialize cuGraph communications
cugraph_comms_init(
    rank=global_rank,
    world_size=world_size,
    uid=cugraph_id[0],
    device=local_rank
)

# Initialize WholeMemory
wm_init(global_rank, world_size, local_rank, torch.cuda.device_count())

# ... training code ...

# Cleanup
wm_finalize()
from pylibcugraph.comms import cugraph_comms_shutdown
cugraph_comms_shutdown()
```

--------------------------------

### Manage Device Memory with rmm::device_buffer

Source: https://github.com/rapidsai/cugraph-gnn/blob/main/cpp/docs/DEVELOPER_GUIDE.md

Demonstrates how to allocate, copy, and move untyped device memory using rmm::device_buffer. It highlights the use of custom memory resources and stream-ordered operations.

```C++
// Allocates at least 100 bytes of uninitialized device memory
// using the specified resource and stream
rmm::device_buffer buff(100, stream, mr);
void * raw_data = buff.data(); // Raw pointer to underlying device memory

// Deep copies `buff` into `copy` on `stream`
rmm::device_buffer copy(buff, stream);

// Moves contents of `buff` into `moved_to`
rmm::device_buffer moved_to(std::move(buff));

custom_memory_resource *mr...;
// Allocates 100 bytes from the custom_memory_resource
rmm::device_buffer custom_buff(100, mr, stream);
```

--------------------------------

### NeighborLoader: GPU-Accelerated Node Sampling

Source: https://context7.com/rapidsai/cugraph-gnn/llms.txt

Implements GraphSAGE-style neighbor sampling for node classification, offering a GPU-accelerated alternative to PyG's NeighborLoader. Requires GraphStore and FeatureStore setup.

```python
import torch
from cugraph_pyg.data import GraphStore, FeatureStore
from cugraph_pyg.loader import NeighborLoader

# Setup data stores
graph_store = GraphStore()
feature_store = FeatureStore()

# Load graph and features (partitioned across workers)
edge_index = torch.load(f"edges_rank{rank}.pt")
graph_store[("node", "rel", "node"), "coo", False, (num_nodes, num_nodes)] = edge_index

features = torch.load(f"features_rank{rank}.pt")
feature_store["node", "x", None] = features
feature_store["node", "y", None] = torch.load(f"labels_rank{rank}.pt")

# Create data tuple
data = (feature_store, graph_store)

# Define input nodes for training
train_nodes = torch.arange(0, 10000, dtype=torch.int64)
```

--------------------------------

### Configure WholeGraph Benchmark Executables with CMake

Source: https://github.com/rapidsai/cugraph-gnn/blob/main/cpp/bench/CMakeLists.txt

The 'ConfigureBench' function automates the creation of benchmark executables by setting include directories, linking required libraries like raft and rmm, and applying specific compiler properties. It supports conditional compilation for NVSHMEM and manages installation paths for testing components.

```cmake
function(ConfigureBench)

  set(options OPTIONAL)
  set(oneValueArgs NAME)
  set(multiValueArgs PATH TARGETS CONFIGURATIONS)
  cmake_parse_arguments(ConfigureBench "${options}" "${oneValueArgs}" "${multiValueArgs}" ${ARGN})

  set(BENCH_NAME ${ConfigureBench_NAME})

  add_executable(${BENCH_NAME} ${ConfigureBench_PATH})

  target_include_directories(
    ${BENCH_NAME} PRIVATE "$<BUILD_INTERFACE:${WHOLEGRAPH_SOURCE_DIR}>/src"
  )
  target_link_libraries(${BENCH_NAME} PRIVATE wholegraph raft::raft rmm::rmm pthread)
  if(BUILD_WITH_NVSHMEM)
    target_compile_definitions(${BENCH_NAME} PRIVATE WITH_NVSHMEM_SUPPORT)
  endif()
  set_target_properties(
    ${BENCH_NAME}
    PROPERTIES
               INSTALL_RPATH "$ORIGIN/../../../lib"
               CXX_STANDARD 17
               CXX_STANDARD_REQUIRED ON
               CUDA_ARCHITECTURES "${CMAKE_CUDA_ARCHITECTURES}"
               POSITION_INDEPENDENT_CODE ON
               RUNTIME_OUTPUT_DIRECTORY "$<BUILD_INTERFACE:${WHOLEGRAPH_BINARY_DIR}/gbench>"
               INTERFACE_POSITION_INDEPENDENT_CODE ON
  )
  target_compile_options(
    ${BENCH_NAME} PUBLIC $<$<COMPILE_LANG_AND_ID:CXX,GNU,Clang>:-Wall -Werror
                         -Wno-error=deprecated-declarations>
  )

  install(
    TARGETS ${BENCH_NAME}
    COMPONENT testing
    DESTINATION bin/gbench/libwholegraph
    EXCLUDE_FROM_ALL
  )

endfunction()
```

--------------------------------

### Enforce Compile-Time Conditions with static_assert (C++)

Source: https://github.com/rapidsai/cugraph-gnn/blob/main/cpp/docs/DEVELOPER_GUIDE.md

The `static_assert` keyword enforces conditions at compile time. This example demonstrates ensuring a template type `T` meets the `std::is_trivial` requirement before proceeding.

```c++
template <typename T>
void trivial_types_only(T t){
   static_assert(std::is_trivial<T>::value, "This function requires a trivial type.");
...
```

--------------------------------

### C++ Template Function and Class with Member Initialization

Source: https://github.com/rapidsai/cugraph-gnn/blob/main/cpp/docs/DEVELOPER_GUIDE.md

Demonstrates a C++ template function accepting a graph object and a template class with private member initialization. This follows snake_case for function and class names, and uses underscore suffix for private members and curly braces for member initialization.

```cpp
template <typename graph_t>
void algorithm_function(graph_t const &g)
{
  ...
}

template <typename vertex_t>
class utility_class
{
  ...
 private:
  vertex_t num_vertices_{};
}
```

--------------------------------

### Manage Distributed Tensors and Embeddings

Source: https://context7.com/rapidsai/cugraph-gnn/llms.txt

Provides examples for initializing, loading, and performing index operations on DistTensor and DistEmbedding objects. These classes facilitate distributed data handling across CPU/GPU backends like NCCL or NVSHMEM.

```python
from cugraph_pyg.tensor import DistTensor, DistEmbedding

# Create DistTensor from PyTorch tensor
host_tensor = torch.randn(1000000, 128, dtype=torch.float32)
dist_tensor = DistTensor.from_tensor(
    tensor=host_tensor,
    device="cpu",
    backend="nccl",
)

# Index operations (distributed gather/scatter)
indices = torch.tensor([0, 100, 1000], dtype=torch.int64, device="cuda")
values = dist_tensor[indices]  # gather
dist_tensor[indices] = torch.randn(3, 128, dtype=torch.float32, device="cuda") # scatter

# DistEmbedding initialization
dist_embedding = DistEmbedding(
    shape=[10000000, 128],
    dtype=torch.float32,
    device="cpu",
    backend="nccl",
    name="node_embeddings"
)
```

--------------------------------

### Enforce Runtime Conditions with wholegraph_EXPECTS (C++)

Source: https://github.com/rapidsai/cugraph-gnn/blob/main/cpp/docs/DEVELOPER_GUIDE.md

The `wholegraph_EXPECTS` macro enforces runtime conditions. If the condition is false, it throws a `wholegraph::logic_error` with a provided message. This is useful for validating inputs and internal states.

```c++
wholegraph_EXPECTS(lhs.type() == rhs.type(), "Column type mismatch");
```

--------------------------------

### NeighborLoader with Temporal Sampling in PyTorch

Source: https://context7.com/rapidsai/cugraph-gnn/llms.txt

Demonstrates setting up NeighborLoader for temporal graphs, enabling time-aware neighbor sampling. This involves adding edge features that include timestamps to the feature store. The loader can then be configured to consider temporal aspects during sampling.

```python
import torch
from cugraph_pyg.data import GraphStore, FeatureStore
from cugraph_pyg.loader import NeighborLoader

graph_store = GraphStore()
feature_store = FeatureStore()

# Add edge features including timestamps (example placeholders)
# feature_store["node", "x", None] = node_features
# feature_store["node", "time", None] = node_timestamps  # per-node timestamps
# feature_store[("node", "rel", "node"), "time", None] = edge_timestamps  # per-edge timestamps

data = (feature_store, graph_store)

# Further configuration for NeighborLoader with temporal sampling would follow,
# specifying temporal constraints or sampling strategies if available in the API.
```

--------------------------------

### Indicate Unreachable Code with wholegraph_FAIL (C++)

Source: https://github.com/rapidsai/cugraph-gnn/blob/main/cpp/docs/DEVELOPER_GUIDE.md

The `wholegraph_FAIL` macro is used to indicate code paths that should never be reached, such as the default case in a switch statement. It throws a `wholegraph::logic_error` with a specified reason.

```c++
wholegraph_FAIL("This code path should not be reached.");
```

--------------------------------

### Initialize CuGraphSampler for PyG

Source: https://github.com/rapidsai/cugraph-gnn/blob/main/readme_pages/cugraph_pyg.md

This code illustrates the initialization of CuGraphSampler, a component that leverages cuGraph for efficient graph sampling within the PyTorch Geometric framework. It takes the feature and graph stores as input and configures sampling parameters like shuffle, number of neighbors, and batch size.

```python
from cugraph_pyg import CuGraphSampler

# Assuming feature_store and graph_store are already created
sampler = CuGraphSampler(
    data=(feature_store, graph_store),
    shuffle=True,
    num_neighbors=[10, 25],
    batch_size=50,
)
```

--------------------------------

### Build Options Configuration

Source: https://github.com/rapidsai/cugraph-gnn/blob/main/python/libwholegraph/CMakeLists.txt

Sets build options for tests and benchmarks to OFF, and enables the use of NCCL runtime wheel. These settings control which components are built and how external libraries are handled.

```cmake
SET(BUILD_TESTS OFF)
SET(BUILD_BENCHMARKS OFF)
SET(USE_NCCL_RUNTIME_WHEEL ON)
```

--------------------------------

### Create NeighborLoader for Homogeneous Graphs in PyTorch

Source: https://context7.com/rapidsai/cugraph-gnn/llms.txt

Demonstrates how to initialize and use the NeighborLoader for homogeneous graphs in a PyTorch training loop. It specifies neighbor sampling fanouts, input nodes, batch size, and shuffling. The loader is used to iterate through batches for model training.

```python
import torch
from cugraph_pyg.data import GraphStore, FeatureStore
from cugraph_pyg.loader import NeighborLoader

# Assuming 'data' is pre-configured with FeatureStore and GraphStore
# Assuming 'train_nodes' are defined

loader = NeighborLoader(
    data,
    num_neighbors=[15, 10, 5],  # fanout per layer: 15 neighbors at hop 1, etc.
    input_nodes=train_nodes,
    batch_size=1024,
    shuffle=True,
    drop_last=True,
    replace=False,              # sample without replacement
    local_seeds_per_call=20000, # tune for memory usage
)

# Training loop example (assuming 'model', 'F', and 'optimizer' are defined)
# model = GNN(in_channels=128, hidden_channels=256, num_classes=10)
# optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

# for epoch in range(10):
#     for batch in loader:
#         batch = batch.to("cuda")
#         optimizer.zero_grad()

#         out = model(batch.x, batch.edge_index)
#         loss = F.cross_entropy(out[:batch.batch_size], batch.y[:batch.batch_size])

#         loss.backward()
#         optimizer.step()

#         print(f"Batch size: {batch.batch_size}, Loss: {loss.item():.4f}")
```

--------------------------------

### Check CUDA API Calls with CUDA_TRY (C++)

Source: https://github.com/rapidsai/cugraph-gnn/blob/main/cpp/docs/DEVELOPER_GUIDE.md

The `CUDA_TRY` macro checks the return value of CUDA runtime API functions. If an error occurs (return value is not `cudaSuccess`), it throws a `wholegraph::cuda_error` exception containing the error description.

```c++
CUDA_TRY( cudaMemcpy(&dst, &src, num_bytes) );
```

--------------------------------

### Clone cuGraph Repository

Source: https://github.com/rapidsai/cugraph-gnn/blob/main/readme_pages/CONTRIBUTING.md

This command clones your forked version of the cuGraph repository to your local machine. Ensure you replace '<YOUR GITHUB NAME>' with your actual GitHub username.

```git
git clone https://github.com/<YOUR GITHUB NAME>/cugraph.git
```

--------------------------------

### NeighborLoader for Heterogeneous Graphs in PyTorch

Source: https://context7.com/rapidsai/cugraph-gnn/llms.txt

Illustrates the use of NeighborLoader with heterogeneous graphs, supporting multiple node and edge types. It shows how to configure stores, add different types of edges and features, and define per-edge-type fanout sampling. The loader outputs a HeteroData object for processing.

```python
import torch
from cugraph_pyg.data import GraphStore, FeatureStore
from cugraph_pyg.loader import NeighborLoader

# Create stores
graph_store = GraphStore()
feature_store = FeatureStore()

# Add heterogeneous edges (example placeholders)
# graph_store[("paper", "cites", "paper"), "coo", False, (100000, 100000)] = paper_cites_paper_edges
# graph_store[("author", "writes", "paper"), "coo", False, (50000, 100000)] = author_writes_paper_edges
# graph_store[("paper", "written_by", "author"), "coo", False, (100000, 50000)] = paper_written_by_author_edges

# Add heterogeneous features (example placeholders)
# feature_store["paper", "x", None] = paper_features   # (100000, 128)
# feature_store["author", "x", None] = author_features # (50000, 64)
# feature_store["paper", "y", None] = paper_labels

data = (feature_store, graph_store)

# Per-edge-type fanout configuration
num_neighbors = {
    ("paper", "cites", "paper"): [10, 5],
    ("author", "writes", "paper"): [5, 3],
    ("paper", "written_by", "author"): [5, 3],
}

# Create loader with heterogeneous input (example placeholders)
# loader = NeighborLoader(
#     data,
#     num_neighbors=num_neighbors,
#     input_nodes=("paper", train_paper_indices),  # (node_type, indices)
#     batch_size=512,
#     shuffle=True,
# )

# for batch in loader:
#     # batch is a HeteroData object
#     # Assuming 'model' and 'F' are defined for heterogeneous graphs
#     # paper_out = model(batch.x_dict, batch.edge_index_dict)
#     # loss = F.cross_entropy(
#     #     paper_out["paper"][:batch["paper"].batch_size],
#     #     batch["paper"].y[:batch["paper"].batch_size]
#     # )

```

--------------------------------

### WholeMemory Initialization

Source: https://context7.com/rapidsai/cugraph-gnn/llms.txt

Initializes the WholeGraph distributed memory system for multi-GPU training, including PyTorch distributed and cuGraph communications.

```APIDOC
## WholeMemory Initialization

### Description
Initialize the WholeGraph distributed memory system for multi-GPU training.

### Method
Initialization code block

### Endpoint
N/A

### Parameters
N/A

### Request Example
```python
import torch
import torch.distributed as dist
from pylibwholegraph.torch.initialize import init as wm_init, finalize as wm_finalize
from pylibcugraph.comms import cugraph_comms_init, cugraph_comms_create_unique_id
import os

# Initialize PyTorch distributed
dist.init_process_group("nccl")
world_size = dist.get_world_size()
global_rank = dist.get_rank()
local_rank = int(os.environ["LOCAL_RANK"])

# Set CUDA device
torch.cuda.set_device(local_rank)

# Create cuGraph unique ID (only rank 0)
if global_rank == 0:
    cugraph_id = [cugraph_comms_create_unique_id()]
else:
    cugraph_id = [None]
dist.broadcast_object_list(cugraph_id, src=0)

# Initialize cuGraph communications
cugraph_comms_init(
    rank=global_rank,
    world_size=world_size,
    uid=cugraph_id[0],
    device=local_rank
)

# Initialize WholeMemory
wm_init(global_rank, world_size, local_rank, torch.cuda.device_count())

# ... training code ...

# Cleanup
wm_finalize()
from pylibcugraph.comms import cugraph_comms_shutdown
cugraph_comms_shutdown()
```

### Response
N/A

### Response Example
N/A
```

--------------------------------

### Register Project Tests

Source: https://github.com/rapidsai/cugraph-gnn/blob/main/cpp/tests/CMakeLists.txt

This snippet demonstrates the usage of the ConfigureTest wrapper to register various unit tests for the project, including parallel utilities, memory communicators, and graph sampling operations.

```cmake
function(ConfigureTest TEST_NAME)
  ConfigureTestInternal(${TEST_NAME} ${ARGN})
endfunction()

ConfigureTest(PARALLEL_UTILS_TEST parallel_utils_tests.cpp)
ConfigureTest(WHOLEMEMORY_TENSOR_TEST wholememory/wholememory_tensor_tests.cpp)
ConfigureTest(
  WHOLEGRAPH_CSR_UNWEIGHTED_SAMPLE_WITHOUT_REPLACEMENT_TEST
  wholegraph_ops/wholegraph_csr_unweighted_sample_without_replacement_tests.cu
  wholegraph_ops/graph_sampling_test_utils.cu
)
```

--------------------------------

### Temporal Graph Neighbor Sampling with cuGraph-GNN NeighborLoader

Source: https://context7.com/rapidsai/cugraph-gnn/llms.txt

Demonstrates how to use the NeighborLoader from cuGraph-GNN to sample temporal neighbors for graph neural network training. It specifies input nodes, timestamps, and temporal sampling strategies.

```python
import torch
from torch_geometric.loader import NeighborLoader

# Assuming 'data' is a PyG Data object with node features, edge_index, and time attributes
# Assuming 'model' is a PyG MessagePassing model

# Input nodes with associated timestamps
input_nodes = torch.arange(1000)
input_time = torch.tensor([100.0] * 1000)  # sample neighbors before time 100

loader = NeighborLoader(
    data,
    num_neighbors=[10, 5],
    input_nodes=input_nodes,
    input_time=input_time,
    time_attr="time",  # name of time attribute in feature store
    temporal_strategy="uniform",
    temporal_comparison="monotonically_decreasing",  # sample older neighbors
    batch_size=256,
)

for batch in loader:
    # Only neighbors with timestamps < seed timestamp are sampled
    out = model(batch.x, batch.edge_index)
```

--------------------------------

### Distributed Multi-GPU GCN Training with cuGraph-PyG

Source: https://context7.com/rapidsai/cugraph-gnn/llms.txt

A complete implementation of distributed GCN training using PyTorch, cuGraph-PyG, and WholeGraph. The script handles distributed environment initialization, graph data partitioning, and neighbor sampling for efficient training on large datasets like OGB-Products.

```python
import os
import torch
import torch.distributed as dist
import torch.nn.functional as F
from torch.nn.parallel import DistributedDataParallel

from pylibwholegraph.torch.initialize import init as wm_init, finalize as wm_finalize
from pylibcugraph.comms import cugraph_comms_init, cugraph_comms_create_unique_id

from cugraph_pyg.data import GraphStore, FeatureStore
from cugraph_pyg.loader import NeighborLoader

import torch_geometric

def setup_distributed():
    """Initialize distributed training environment."""
    dist.init_process_group("nccl")

    world_size = dist.get_world_size()
    global_rank = dist.get_rank()
    local_rank = int(os.environ["LOCAL_RANK"])

    torch.cuda.set_device(local_rank)

    # Create cuGraph unique ID
    if global_rank == 0:
        cugraph_id = [cugraph_comms_create_unique_id()]
    else:
        cugraph_id = [None]
    dist.broadcast_object_list(cugraph_id, src=0)

    # Initialize cuGraph and WholeMemory
    cugraph_comms_init(
        rank=global_rank,
        world_size=world_size,
        uid=cugraph_id[0],
        device=local_rank
    )
    wm_init(global_rank, world_size, local_rank, torch.cuda.device_count())

    return global_rank, local_rank, world_size

def load_data(rank, world_size, data_dir):
    """Load partitioned graph data."""
    graph_store = GraphStore()
    feature_store = FeatureStore()

    # Load this rank's partition
    edge_index = torch.load(f"{data_dir}/edges_rank{rank}.pt")
    features = torch.load(f"{data_dir}/features_rank{rank}.pt")
    labels = torch.load(f"{data_dir}/labels_rank{rank}.pt")

    num_nodes = 2449029  # OGB-Products

    graph_store[
        ("node", "rel", "node"), "coo", False, (num_nodes, num_nodes)
    ] = edge_index

    feature_store["node", "x", None] = features
    feature_store["node", "y", None] = labels

    # Load split indices
    train_idx = torch.load(f"{data_dir}/train_idx_rank{rank}.pt")
    val_idx = torch.load(f"{data_dir}/val_idx_rank{rank}.pt")
    test_idx = torch.load(f"{data_dir}/test_idx_rank{rank}.pt")

    return (feature_store, graph_store), train_idx, val_idx, test_idx

def train(rank, local_rank, world_size, epochs=10):
    """Main training function."""
    device = torch.device(f"cuda:{local_rank}")

    # Load data
    data, train_idx, val_idx, test_idx = load_data(rank, world_size, "/data/ogb")

    # Create model
    model = torch_geometric.nn.models.GCN(
        in_channels=100,
        hidden_channels=256,
        num_layers=3,
        out_channels=47,
    ).to(device)
    model = DistributedDataParallel(model, device_ids=[local_rank])

    optimizer = torch.optim.Adam(model.parameters(), lr=0.01)

    # Create loaders
    train_loader = NeighborLoader(
        data,
        num_neighbors=[15, 10, 5],
        input_nodes=train_idx.cuda(),
        batch_size=1024,
        shuffle=True,
        drop_last=True,
    )

    val_loader = NeighborLoader(
        data,
        num_neighbors=[15, 10, 5],
        input_nodes=val_idx.cuda(),
        batch_size=1024,
        drop_last=True,
    )

    # Training loop
    for epoch in range(epochs):
        model.train()
        total_loss = 0

        for batch in train_loader:
            batch = batch.to(device)
            optimizer.zero_grad()

            out = model(batch.x, batch.edge_index)
            loss = F.cross_entropy(
                out[:batch.batch_size],
                batch.y[:batch.batch_size].view(-1)
            )

            loss.backward()
            optimizer.step()
            total_loss += loss.item()

        # Validation
        model.eval()
        correct = total = 0

        with torch.no_grad():
            for batch in val_loader:
                batch = batch.to(device)
                out = model(batch.x, batch.edge_index)[:batch.batch_size]
                pred = out.argmax(dim=-1)
                correct += (pred == batch.y[:batch.batch_size].view(-1)).sum().item()
                total += batch.batch_size

        if rank == 0:
            print(f"Epoch {epoch}: Loss={total_loss:.4f}, Val Acc={100*correct/total:.2f}%")

def main():
    global_rank, local_rank, world_size = setup_distributed()

    try:
        train(global_rank, local_rank, world_size, epochs=10)
    finally:
        wm_finalize()
        from pylibcugraph.comms import cugraph_comms_shutdown
        cugraph_comms_shutdown()

if __name__ == "__main__":
    main()
```

--------------------------------

### GraphStore: Distributed Graph Storage in PyG

Source: https://context7.com/rapidsai/cugraph-gnn/llms.txt

Implements a cuGraph-backed PyG GraphStore for distributing graph structures across workers. Supports lazy graph creation and various deployment configurations. Requires PyTorch and PyTorch distributed initialization.

```python
import torch
import torch.distributed as dist
from cugraph_pyg.data import GraphStore, FeatureStore

# Initialize distributed environment
dist.init_process_group("nccl")

# Create graph store
graph_store = GraphStore()

# Add edge index in COO format
# Format: (src_type, relation_type, dst_type)
edge_index = torch.tensor([[0, 1, 2, 3], [1, 2, 3, 0]], dtype=torch.int64)
num_nodes = 1000

# Put edge index into the store
graph_store[
    ("node", "connects", "node"),  # edge type tuple
    "coo",                          # layout format
    False,                          # is_sorted
    (num_nodes, num_nodes)          # size tuple (num_src, num_dst)
] = edge_index

# For heterogeneous graphs with multiple edge types
graph_store[
    ("user", "rates", "movie"),
    "coo",
    False,
    (10000, 5000)  # 10000 users, 5000 movies
] = user_movie_edges

graph_store[
    ("movie", "belongs_to", "genre"),
    "coo",
    False,
    (5000, 100)  # 5000 movies, 100 genres
] = movie_genre_edges

# Get all edge attributes
edge_attrs = graph_store.get_all_edge_attrs()
for attr in edge_attrs:
    print(f"Edge type: {attr.edge_type}, size: {attr.size}")
```

--------------------------------

### Create and Train WholeMemory Embeddings

Source: https://context7.com/rapidsai/cugraph-gnn/llms.txt

Illustrates the creation, training, and management of distributed embeddings using WholeMemoryEmbedding. It covers initialization, optimizer creation, training steps, PyTorch Module integration, saving/loading, and cache management.

```python
import torch
from pylibwholegraph.torch.embedding import (
    WholeMemoryEmbedding,
    WholeMemoryEmbeddingModule,
    WholeMemoryOptimizer,
    create_embedding,
    create_embedding_from_filelist,
    create_wholememory_optimizer,
    create_builtin_cache_policy,
    destroy_embedding,
)
from pylibwholegraph.torch.comm import get_global_communicator

comm = get_global_communicator()

# Create embedding with optional caching
cache_policy = create_builtin_cache_policy(
    builtin_cache_type="local_node",  # "none", "local_device", "local_node", "all_devices"
    embedding_memory_type="distributed",
    embedding_memory_location="cpu",
    access_type="readonly",
    cache_ratio=0.2,  # cache 20% of embeddings
)

# Create distributed embedding
embedding = create_embedding(
    comm=comm,
    memory_type="distributed",  # "continuous", "chunked", "distributed"
    memory_location="cpu",      # store in host memory (UVA)
    dtype=torch.float32,
    sizes=[10000000, 128],      # 10M embeddings, 128 dimensions
    cache_policy=cache_policy,
    random_init=True,           # Xavier initialization
)

# Access embedding properties
print(f"Embedding shape: {embedding.shape}")

# Gather embeddings (forward pass)
indices = torch.tensor([0, 1, 2, 3], dtype=torch.int64, device="cuda")
output = embedding.gather(indices, is_training=True)
print(f"Output shape: {output.shape}")

# Create optimizer for training
optimizer = create_wholememory_optimizer(
    embeddings=embedding,
    optimizer_type="lazy_adam",  # "sgd", "lazy_adam", "rmsprop", "adagrad"
    param_dict={"beta1": 0.9, "beta2": 0.999, "epsilon": 1e-8},
)

# Training step
loss = some_loss_function(output)
loss.backward()
optimizer.step(lr=0.001)

# PyTorch Module wrapper for easier integration
embedding_module = WholeMemoryEmbeddingModule(embedding)

# Use in PyTorch model
class GNNWithEmbedding(torch.nn.Module):
    def __init__(self, embedding):
        super().__init__()
        self.embedding = WholeMemoryEmbeddingModule(embedding)
        self.conv1 = torch_geometric.nn.GCNConv(128, 256)
        self.conv2 = torch_geometric.nn.GCNConv(256, 10)

    def forward(self, node_ids, edge_index):
        x = self.embedding(node_ids)  # distributed embedding lookup
        x = F.relu(self.conv1(x, edge_index))
        x = self.conv2(x, edge_index)
        return x

# Save/load embeddings
embedding.save("/checkpoints/embedding")
embedding.load("/checkpoints/embedding", part_count=world_size)

# Cache management
embedding.writeback_all_cache()  # write cache back to storage
embedding.drop_all_cache()       # invalidate cache

# Cleanup
destroy_embedding(embedding)
```

--------------------------------

### Initialize CUDA Architectures for LIBWHOLEGRAPH

Source: https://github.com/rapidsai/cugraph-gnn/blob/main/python/libwholegraph/CMakeLists.txt

Initializes CUDA architectures for the LIBWHOLEGRAPH target. This function is part of the rapids-cuda module and prepares the build environment for CUDA compilation.

```cmake
rapids_cuda_init_architectures(LIBWHOLEGRAPH)
```

--------------------------------

### Save and Load WholeMemory Tensors

Source: https://context7.com/rapidsai/cugraph-gnn/llms.txt

Demonstrates how to save and load WholeMemory tensors to and from files using a file prefix. This is useful for persisting and retrieving distributed tensor data across different sessions or ranks.

```python
# Save tensor to files (each rank saves its partition)
wm_tensor.to_file_prefix("/data/embeddings")

# Load tensor from files
wm_tensor.from_file_prefix("/data/embeddings", part_count=world_size)

# Create tensor directly from file list
wm_tensor_from_files = create_wholememory_tensor_from_filelist(
    comm=comm,
    memory_type="chunked",
    memory_location="cpu",
    filelist=["/data/part0.bin", "/data/part1.bin"],
    dtype=torch.float32,
    last_dim_size=128,
)

# Cleanup
destroy_wholememory_tensor(wm_tensor)
```

--------------------------------

### Build Cython Modules with CMake

Source: https://github.com/rapidsai/cugraph-gnn/blob/main/python/pylibwholegraph/pylibwholegraph/binding/CMakeLists.txt

This snippet demonstrates how to define Cython source files and link them against the WholeGraph library using the RAPIDS CMake utilities. It automates the creation of Cython modules by specifying the source files and associated target dependencies.

```cmake
set(cython_sources wholememory_binding.pyx)
set(linked_libraries wholegraph::wholegraph)

rapids_cython_create_modules(
  CXX
  SOURCE_FILES "${cython_sources}"
  LINKED_LIBRARIES "${linked_libraries}" ASSOCIATED_TARGETS wholegraph MODULE_PREFIX wholegraphcif
)
```

--------------------------------

### WholeGraph Communicator Management (pylibwholegraph)

Source: https://context7.com/rapidsai/cugraph-gnn/llms.txt

Demonstrates creating and managing communicators for distributed memory operations within the WholeGraph system. It covers global, local node, local device, and custom group communicators.

```python
from pylibwholegraph.torch.comm import (
    create_group_communicator,
    get_global_communicator,
    get_local_node_communicator,
    get_local_device_communicator,
    split_communicator,
    set_world_info,
)
import torch

# Assuming global_rank, world_size, local_rank are already defined
# Set world information
set_world_info(
    world_rank=global_rank,
    world_size=world_size,
    local_rank=local_rank,
    local_size=torch.cuda.device_count()
)

# Get global communicator (all GPUs across all nodes)
global_comm = get_global_communicator(distributed_backend="nccl")
print(f"Global rank: {global_comm.get_rank()}, size: {global_comm.get_size()}")

# Get local node communicator (GPUs within same node)
local_comm = get_local_node_communicator()
print(f"Local rank: {local_comm.get_rank()}, size: {local_comm.get_size()}")

# Get local device communicator (single GPU)
device_comm = get_local_device_communicator()

# Create custom group communicator
# Example: 24 ranks with group_size=4, stride=2 creates groups:
# [0,2,4,6], [1,3,5,7], [8,10,12,14], ...
custom_comm = create_group_communicator(group_size=4, comm_stride=2)

# Split communicator by color
new_comm = split_communicator(global_comm, color=global_rank % 4, key=global_rank)

# Barrier synchronization
global_comm.barrier()

# Check supported memory configurations
supports_chunked_cuda = global_comm.support_type_location("chunked", "cuda")
print(f"Supports chunked CUDA: {supports_chunked_cuda}")
```

--------------------------------

### Checkout cuGraph Branch

Source: https://github.com/rapidsai/cugraph-gnn/blob/main/readme_pages/CONTRIBUTING.md

This command checks out the specified branch for cuGraph development. Replace 'branch-x.x' with the actual name of the current development branch, which can be found on the main cuGraph GitHub page.

```git
git checkout branch-x.x
```

--------------------------------

### Convert cuGraph Property Graph to PyG Stores

Source: https://github.com/rapidsai/cugraph-gnn/blob/main/readme_pages/cugraph_pyg.md

This snippet demonstrates how to convert a cuGraph PropertyGraph object into PyTorch Geometric compatible FeatureStore and GraphStore objects. It's a crucial step for integrating cuGraph data into PyG's data handling pipeline.

```python
import cugraph
from cugraph_pyg import to_pyg

G = cugraph.PropertyGraph()
# ... populate G with graph data ...
feature_store, graph_store = to_pyg(G)
```

--------------------------------

### Set Runtime Path for NCCL

Source: https://github.com/rapidsai/cugraph-gnn/blob/main/python/libwholegraph/CMakeLists.txt

Appends the runtime library path for NCCL if USE_NCCL_RUNTIME_WHEEL is enabled. This ensures that the dynamic linker can find the necessary NCCL libraries at runtime.

```cmake
if(USE_NCCL_RUNTIME_WHEEL)
  list(APPEND rpaths "$ORIGIN/../../nvidia/nccl/lib")
endif()

set_property(
  TARGET wholegraph
  PROPERTY INSTALL_RPATH ${rpaths}
  APPEND
)
```

--------------------------------

### Initialize CUDA Architectures for PYLIBWHOLEGRAPH

Source: https://github.com/rapidsai/cugraph-gnn/blob/main/python/pylibwholegraph/CMakeLists.txt

Initializes CUDA architectures for the PYLIBWHOLEGRAPH target. This function is part of the rapids-cmake framework and prepares the build environment for CUDA compilation specific to the whole graph Python library.

```cmake
rapids_cuda_init_architectures(PYLIBWHOLEGRAPH)
```

--------------------------------

### Configure Test Executables with CMake

Source: https://github.com/rapidsai/cugraph-gnn/blob/main/cpp/tests/CMakeLists.txt

The ConfigureTestInternal function encapsulates the logic for creating test executables, linking necessary libraries (GTest, raft, rmm), and setting target properties. It includes conditional logic for NVSHMEM support, code coverage instrumentation, and compiler-specific warning suppression.

```cmake
function(ConfigureTestInternal TEST_NAME)
  add_executable(${TEST_NAME} ${ARGN})
  target_include_directories(${TEST_NAME} PRIVATE "$<BUILD_INTERFACE:${WHOLEGRAPH_SOURCE_DIR}>/src")
  target_link_libraries(
    ${TEST_NAME} PRIVATE GTest::gmock GTest::gtest GTest::gmock_main GTest::gtest_main wholegraph
                         raft::raft rmm::rmm pthread
  )
  if(BUILD_WITH_NVSHMEM)
    target_compile_definitions(${TEST_NAME} PRIVATE WITH_NVSHMEM_SUPPORT)
  endif()

  set_target_properties(
    ${TEST_NAME}
    PROPERTIES POSITION_INDEPENDENT_CODE ON
               RUNTIME_OUTPUT_DIRECTORY "$<BUILD_INTERFACE:${WHOLEGRAPH_BINARY_DIR}/gtests>"
               CUDA_ARCHITECTURES "${CMAKE_CUDA_ARCHITECTURES}"
               INSTALL_RPATH "\$ORIGIN/../../../lib"
  )
  # ... (additional configuration logic)
  add_test(NAME ${TEST_NAME} COMMAND ${TEST_NAME})
endfunction()
```

--------------------------------

### Create and Use CSR GraphStructure

Source: https://context7.com/rapidsai/cugraph-gnn/llms.txt

Shows how to create a GraphStructure object representing a graph in Compressed Sparse Row (CSR) format. This includes creating WholeMemory tensors for row pointers and column indices, setting edge attributes, and performing unweighted sampling.

```python
import torch
from pylibwholegraph.torch.graph_structure import GraphStructure
from pylibwholegraph.torch.tensor import create_wholememory_tensor
from pylibwholegraph.torch.comm import get_global_communicator

comm = get_global_communicator()

# Create CSR row pointers and column indices as WholeMemory tensors
csr_row_ptr = create_wholememory_tensor(
    comm=comm,
    memory_type="chunked",
    memory_location="cuda",
    sizes=[num_nodes + 1],
    dtype=torch.int64,
    strides=[1],
)

csr_col_ind = create_wholememory_tensor(
    comm=comm,
    memory_type="chunked",
    memory_location="cuda",
    sizes=[num_edges],
    dtype=torch.int64,
    strides=[1],
)

# Initialize CSR data from local partitions
local_row_ptr, _ = csr_row_ptr.get_local_tensor()
local_col_ind, _ = csr_col_ind.get_local_tensor()
# ... fill with graph data ...

# Create graph structure
graph = GraphStructure()
graph.set_csr_graph(csr_row_ptr, csr_col_ind)

# Add edge attributes (e.g., weights)
edge_weights = create_wholememory_tensor(
    comm=comm,
    memory_type="chunked",
    memory_location="cuda",
    sizes=[num_edges],
    dtype=torch.float32,
    strides=[1],
)
graph.set_edge_attribute("weight", edge_weights)

# Unweighted sampling without replacement
center_nodes = torch.tensor([0, 1, 2, 3], dtype=torch.int64, device="cuda")
csr_row_ptr_out, sampled_nodes = graph.unweighted_sample_without_replacement_one_hop(
    center_nodes_tensor=center_nodes,
    max_sample_count=10,
    random_seed=42,
)

```

--------------------------------

### FeatureStore: Distributed Feature Storage

Source: https://context7.com/rapidsai/cugraph-gnn/llms.txt

Provides WholeGraph-backed distributed feature storage that avoids data replication. Features are automatically distributed and accessible from any worker. Supports CPU or CUDA locations.

```python
import torch
from cugraph_pyg.data import FeatureStore

# Create feature store with specified location
feature_store = FeatureStore(location="cpu")  # or "cuda"

# Store node features (each worker provides its partition)
# Features are automatically distributed across all workers
node_features = torch.randn(1000, 128, dtype=torch.float32)  # local partition
feature_store["node", "x", None] = node_features

# Store node labels
node_labels = torch.randint(0, 10, (1000,), dtype=torch.int64)
feature_store["node", "y", None] = node_labels

# For heterogeneous graphs
user_features = torch.randn(5000, 64, dtype=torch.float32)
feature_store["user", "x", None] = user_features

movie_features = torch.randn(2000, 128, dtype=torch.float32)
feature_store["movie", "x", None] = movie_features

# Retrieve features by index (works across all workers)
indices = torch.tensor([0, 5, 10, 15], dtype=torch.int64, device="cuda")
retrieved_features = feature_store["node", "x", None][indices]

# Get all tensor attributes
tensor_attrs = feature_store.get_all_tensor_attrs()
for attr in tensor_attrs:
    print(f"Group: {attr.group_name}, Attr: {attr.attr_name}")
```

--------------------------------

### Perform Graph Sampling Operations

Source: https://context7.com/rapidsai/cugraph-gnn/llms.txt

Demonstrates how to execute weighted one-hop sampling and multi-layer GraphSAGE-style sampling on graph structures. These methods return node IDs and edge indices necessary for downstream GNN computations.

```python
# Weighted sampling with edge weights
csr_row_ptr_out, sampled_nodes, center_lids, edge_ids = graph.weighted_sample_without_replacement_one_hop(
    weight_name="weight",
    center_nodes_tensor=center_nodes,
    max_sample_count=10,
    random_seed=42,
    need_center_local_output=True,
    need_edge_output=True,
)

# Multi-layer sampling (GraphSAGE style)
seed_nodes = torch.tensor([100, 200, 300], dtype=torch.int64, device="cuda")
target_gids, edge_indices, csr_row_ptrs, csr_col_inds = graph.multilayer_sample_without_replacement(
    node_ids=seed_nodes,
    max_neighbors=[15, 10, 5],  # fanout per layer
    weight_name=None,          # unweighted
)
```

--------------------------------

### WholeGraph Distributed Tensor Operations (pylibwholegraph)

Source: https://context7.com/rapidsai/cugraph-gnn/llms.txt

Shows how to create, manipulate, and perform operations on distributed tensors using WholeMemoryTensor in WholeGraph. Supports creation from scratch, file lists, and operations like gather and scatter.

```python
import torch
from pylibwholegraph.torch.tensor import (
    WholeMemoryTensor,
    create_wholememory_tensor,
    create_wholememory_tensor_from_filelist,
    destroy_wholememory_tensor,
)
from pylibwholegraph.torch.comm import get_global_communicator

comm = get_global_communicator()

# Create empty distributed tensor
wm_tensor = create_wholememory_tensor(
    comm=comm,
    memory_type="chunked",    # "continuous", "chunked", or "distributed"
    memory_location="cuda",   # "cpu" or "cuda"
    sizes=[1000000, 128],     # total shape across all ranks
    dtype=torch.float32,
    strides=[128, 1],         # row-major strides
)

# Access tensor properties
print(f"Shape: {wm_tensor.shape}")
print(f"Dtype: {wm_tensor.dtype}")
print(f"Dimensions: {wm_tensor.dim()}")

# Get local tensor partition
local_tensor, offset = wm_tensor.get_local_tensor()
print(f"Local shape: {local_tensor.shape}, offset: {offset}")

# Initialize local partition with data
torch.nn.init.xavier_uniform_(local_tensor)

# Gather operation (retrieve by indices from any rank)
indices = torch.tensor([0, 100, 1000, 5000], dtype=torch.int64, device="cuda")
gathered = wm_tensor.gather(indices)  # returns tensor on CUDA
print(f"Gathered shape: {gathered.shape}")

# Scatter operation (write by indices)
new_values = torch.randn(4, 128, device="cuda")
wm_tensor.scatter(new_values, indices)

# Get sub-tensor
sub_tensor = wm_tensor.get_sub_tensor(starts=[0, 0], ends=[1000, 64])

# Destroy the tensor to free memory
destroy_wholememory_tensor(wm_tensor)
```

--------------------------------

### Project Definition and Package Finding for PYLIBWHOLEGRAPH

Source: https://github.com/rapidsai/cugraph-gnn/blob/main/python/pylibwholegraph/CMakeLists.txt

Defines the project 'PYLIBWHOLEGRAPH' with its version and languages (CXX, CUDA). It then finds the 'wholegraph' package, ensuring its availability for the build. This is crucial for integrating the 'wholegraph' library with cugraph-gnn.

```cmake
project(
  PYLIBWHOLEGRAPH
  VERSION "${RAPIDS_VERSION}"
  LANGUAGES CXX CUDA
)

find_package(wholegraph "${RAPIDS_VERSION}" REQUIRED)
```

--------------------------------

### Add Subdirectory for libwholegraph

Source: https://github.com/rapidsai/cugraph-gnn/blob/main/python/libwholegraph/CMakeLists.txt

Adds the '../../cpp/' directory as a subdirectory for building the libwholegraph library. This command integrates external source code into the main build process.

```cmake
add_subdirectory(../../cpp/ libwholegraph)
```

--------------------------------

### Add Subdirectory for Python Bindings

Source: https://github.com/rapidsai/cugraph-gnn/blob/main/python/pylibwholegraph/CMakeLists.txt

Adds the 'pylibwholegraph/binding' subdirectory to the build. This typically contains the source files for creating Python bindings, allowing Python code to interact with the underlying C++/CUDA implementations.

```cmake
add_subdirectory(pylibwholegraph/binding)
```