def predict(x: np.ndarray) -> np.ndarray:
            """
            Predict extracted model.

            :param x: Samples of input data of shape `(num_samples, num_features)`.
            :return: Predictions with the extracted model of shape `(num_samples, num_classes)`.
            """
            layer_0 = np.maximum(np.matmul(self.w_0.T, x.T) + self.b_0, 0.0)  # type: ignore
            layer_1 = np.matmul(self.w_1.T, layer_0) + self.b_1  # type: ignore
            return layer_1.T

            Predict extracted model.

            :param x: Samples of input data of shape `(num_samples, num_features)`.
            :return: Predictions with the extracted model of shape `(num_samples, num_classes)`.
            
art/attacks/extraction/functionally_equivalent_extraction.py