def _saliency_map(self, x: np.ndarray, target: np.ndarray | int, search_space: np.ndarray) -> np.ndarray:
        """
        Compute the saliency map of `x`. Return the top 2 coefficients in `search_space` that maximize / minimize
        the saliency map.

        :param x: A batch of input samples.
        :param target: Target class for `x`.
        :param search_space: The set of valid pairs of feature indices to search.
        :return: The top 2 coefficients in `search_space` that maximize / minimize the saliency map.
        """
        grads = self.estimator.class_gradient(x, label=target)
        grads = np.reshape(grads, (-1, self._nb_features))

        # Remove gradients for already used features
        used_features = 1 - search_space
        coeff = 2 * int(self.theta > 0) - 1
        grads[used_features == 1] = -np.inf * coeff

        if self.theta > 0:
            ind = np.argpartition(grads, -2, axis=1)[:, -2:]
        else:  # pragma: no cover
            ind = np.argpartition(-grads, -2, axis=1)[:, -2:]

        return ind

        Compute the saliency map of `x`. Return the top 2 coefficients in `search_space` that maximize / minimize
        the saliency map.

        :param x: A batch of input samples.
        :param target: Target class for `x`.
        :param search_space: The set of valid pairs of feature indices to search.
        :return: The top 2 coefficients in `search_space` that maximize / minimize the saliency map.
        
art/attacks/evasion/saliency_map.py