def query(self, keys):
        """Queries the GPU cache.

        Parameters
        ----------
        keys : Tensor
            The keys to query the GPU graph cache with.

        Returns
        -------
        tuple(Tensor, func)
            A tuple containing (missing_keys, replace_fn) where replace_fn is a
            function that should be called with the graph structure
            corresponding to the missing keys. Its arguments are
            (Tensor, list(Tensor)), where the first tensor is the missing indptr
            and the second list is the missing edge tensors.
        """
        self.total_queries += keys.shape[0]
        (
            index,
            position,
            num_hit,
            num_threshold,
        ) = self._cache.query(keys)
        self.total_miss += keys.shape[0] - num_hit

        def replace_functional(missing_indptr, missing_edge_tensors):
            return self._cache.replace(
                keys,
                index,
                position,
                num_hit,
                num_threshold,
                missing_indptr,
                missing_edge_tensors,
            )

        return keys[index[num_hit:]], replace_functional
Queries the GPU cache.

        Parameters
        ----------
        keys : Tensor
            The keys to query the GPU graph cache with.

        Returns
        -------
        tuple(Tensor, func)
            A tuple containing (missing_keys, replace_fn) where replace_fn is a
            function that should be called with the graph structure
            corresponding to the missing keys. Its arguments are
            (Tensor, list(Tensor)), where the first tensor is the missing indptr
            and the second list is the missing edge tensors.
        
python/dgl/graphbolt/impl/gpu_graph_cache.py