Method compute_blocks

python/dgl/graphbolt/minibatch.py:183–304

Extracts DGL blocks from `MiniBatch` to construct graphical structures and ID mappings.

(self)

        return self._blocks

    def compute_blocks(self) -> list:
        """Extracts DGL blocks from `MiniBatch` to construct graphical
        structures and ID mappings.
        """
        from dgl.convert import create_block, EID, NID

        is_heterogeneous = isinstance(
            self.sampled_subgraphs[0].sampled_csc, Dict
        )

        # Casts to minimum dtype in-place and returns self.
        def cast_to_minimum_dtype(v: CSCFormatBase):
            # Checks if number of vertices and edges fit into an int32.
            dtype = (
                torch.int32
                if max(v.indptr.size(0) - 2, v.indices.size(0))
                <= torch.iinfo(torch.int32).max
                else torch.int64
            )
            v.indptr = v.indptr.to(dtype)
            v.indices = v.indices.to(dtype)
            return v

        blocks = []
        for subgraph in self.sampled_subgraphs:
            original_row_node_ids = subgraph.original_row_node_ids
            assert (
                original_row_node_ids is not None
            ), "Missing `original_row_node_ids` in sampled subgraph."
            original_column_node_ids = subgraph.original_column_node_ids
            assert (
                original_column_node_ids is not None
            ), "Missing `original_column_node_ids` in sampled subgraph."
            if is_heterogeneous:
                node_types = set()
                sampled_csc = {}
                for v in subgraph.sampled_csc.values():
                    cast_to_minimum_dtype(v)
                for etype, v in subgraph.sampled_csc.items():
                    etype_tuple = etype_str_to_tuple(etype)
                    node_types.add(etype_tuple[0])
                    node_types.add(etype_tuple[2])
                    sampled_csc[etype_tuple] = (
                        "csc",
                        (
                            v.indptr,
                            v.indices,
                            torch.arange(
                                0,
                                len(v.indices),
                                device=v.indptr.device,
                                dtype=v.indptr.dtype,
                            ),
                        ),
                    )
                num_src_nodes = {
                    ntype: (
                        original_row_node_ids[ntype].size(0)