Function _get_subgraph_batch_info

tests/python/common/test_batch-graph.py:239–270

Internal function to compute batch information for subgraphs. Parameters ---------- keys : List[str] The node/edge type keys. induced_indices_arr : List[Tensor] The induced node/edge index tensor for all node/edge types. batch_num_objs : Tensor Number of n

(keys, induced_indices_arr, batch_num_objs)

def _get_subgraph_batch_info(keys, induced_indices_arr, batch_num_objs):
    """Internal function to compute batch information for subgraphs.
    Parameters
    ----------
    keys : List[str]
        The node/edge type keys.
    induced_indices_arr : List[Tensor]
        The induced node/edge index tensor for all node/edge types.
    batch_num_objs : Tensor
        Number of nodes/edges for each graph in the original batch.
    Returns
    -------
    Mapping[str, Tensor]
        A dictionary mapping all node/edge type keys to the ``batch_num_objs``
        array of corresponding graph.
    """
    bucket_offset = np.expand_dims(
        np.cumsum(F.asnumpy(batch_num_objs), 0), -1
    )  # (num_bkts, 1)
    ret = {}
    for key, induced_indices in zip(keys, induced_indices_arr):
        # NOTE(Zihao): this implementation is not efficient and we can replace it with
        # binary search in the future.
        induced_indices = np.expand_dims(
            F.asnumpy(induced_indices), 0
        )  # (1, num_nodes)
        new_offset = np.sum((induced_indices < bucket_offset), 1)  # (num_bkts,)
        # start_offset = [0] + [new_offset[i-1] for i in range(1, n_bkts)]
        start_offset = np.concatenate([np.zeros((1,)), new_offset[:-1]], 0)
        new_batch_num_objs = new_offset - start_offset
        ret[key] = F.tensor(new_batch_num_objs, dtype=F.dtype(batch_num_objs))
    return ret


@parametrize_idtype