[GraphBolt][PyG] Add SampledSubgraph.to_pyg.

[mfbalin]

🔨Work Item

IMPORTANT:

• This template is only for dev team to track project progress. For feature request or bug report, please use the corresponding issue templates.
• DO NOT create a new work item if the purpose is to fix an existing issue or feature request. We will directly use the issue in the project tracker.

Project tracker: https://github.com/orgs/dmlc/projects/2

Description

The code in our examples for PyG integration can be made part of GraphBolt. Then we won't have to replicate this code anymore in our PyG examples.

dgl/examples/graphbolt/pyg/node_classification_advanced.py

Lines 78 to 94 in 0331009

	def convert_to_pyg(h, subgraph):
	#####################################################################
	# (HIGHLIGHT) Convert given features to be consumed by a PyG layer.
	#
	# We convert the provided sampled edges in CSC format from GraphBolt and
	# convert to COO via using gb.expand_indptr.
	#####################################################################
	src = subgraph.sampled_csc.indices
	dst = gb.expand_indptr(
	subgraph.sampled_csc.indptr,
	dtype=src.dtype,
	output_size=src.size(0),
	)
	edge_index = torch.stack([src, dst], dim=0).long()
	dst_size = subgraph.sampled_csc.indptr.size(0) - 1
	# h and h[:dst_size] correspond to source and destination features resp.
	return (h, h[:dst_size]), edge_index, (h.size(0), dst_size)

dgl/examples/graphbolt/pyg/hetero/node_classification.py

Lines 77 to 102 in 0331009

	def convert_to_pyg(h, subgraph):
	#####################################################################
	# (HIGHLIGHT) Convert given features to be consumed by a PyG layer.
	#
	# We convert the provided sampled edges in CSC format from GraphBolt and
	# convert to COO via using gb.expand_indptr.
	#####################################################################
	h_dst_dict = {}
	edge_index_dict = {}
	sizes_dict = {}
	for etype, sampled_csc in subgraph.sampled_csc.items():
	src = sampled_csc.indices
	dst = gb.expand_indptr(
	sampled_csc.indptr,
	dtype=src.dtype,
	output_size=src.size(0),
	)
	edge_index = torch.stack([src, dst], dim=0).long()
	dst_size = sampled_csc.indptr.size(0) - 1
	# h and h[:dst_size] correspond to source and destination features resp.
	src_ntype, _, dst_ntype = gb.etype_str_to_tuple(etype)
	h_dst_dict[dst_ntype] = h[dst_ntype][:dst_size]
	edge_index_dict[etype] = edge_index
	sizes_dict[etype] = (h[src_ntype].size(0), dst_size)
	return (h, h_dst_dict), edge_index_dict, sizes_dict