def METHOD_NAME(
    self,
    ids: torch.Tensor,
    node_type: Optional[NodeType] = None,
) -> Tensor:
    r"""Get the partition IDs of node IDs for a specific node type."""
    if self.meta['is_hetero']:
        assert node_type is not None
        return self.node_pb[node_type][ids]
    return self.node_pb[ids]