def nid2partid(self, nids, ntype=DEFAULT_NTYPE):
        """From global node IDs to partition IDs"""
        nids = utils.toindex(nids)
        if ntype == DEFAULT_NTYPE:
            ret = np.searchsorted(
                self._max_node_ids, nids.tonumpy(), side="right"
            )
        else:
            ret = np.searchsorted(
                self._typed_max_node_ids[ntype], nids.tonumpy(), side="right"
            )
        ret = utils.toindex(ret)
        return ret.tousertensor()
From global node IDs to partition IDs
python/dgl/distributed/graph_partition_book.py