def _reward(self, y_output: np.ndarray, y_hat: np.ndarray, n: int) -> float | np.ndarray:
        """
        Compute reward value.

        :param y_output: Output of the victim classifier.
        :param y_hat: Output of the thieved classifier.
        :param n: Current iteration.
        :return: Reward value.
        """
        if self.reward == "cert":
            return self._reward_cert(y_output)
        if self.reward == "div":
            return self._reward_div(y_output, n)
        if self.reward == "loss":
            return self._reward_loss(y_output, y_hat)
        return self._reward_all(y_output, y_hat, n)

        Compute reward value.

        :param y_output: Output of the victim classifier.
        :param y_hat: Output of the thieved classifier.
        :param n: Current iteration.
        :return: Reward value.
        
art/attacks/extraction/knockoff_nets.py