def predict(
        self, x: np.ndarray, batch_size: int = 128, raw: bool = False, max_aggregate: bool = True, **kwargs
    ) -> np.ndarray:
        """
        Perform prediction for a batch of inputs. Aggregation will be performed on the prediction from
        each classifier if max_aggregate is True. Otherwise, the probabilities will be summed instead.
        For logits output set max_aggregate=True, as logits are not comparable between models and should
        not be aggregated using a sum.

        :param x: Input samples.
        :param batch_size: Size of batches.
        :param raw: Return the individual classifier raw outputs (not aggregated).
        :param max_aggregate: Aggregate the predicted classes of each classifier if True. If false, aggregation
               is done using a sum. If raw is true, this arg is ignored
        :return: Array of predictions of shape `(nb_inputs, nb_classes)`, or of shape
                 `(nb_classifiers, nb_inputs, nb_classes)` if `raw=True`.
        """

        if raw:
            return super().predict(x, batch_size=batch_size, raw=True, **kwargs)

        # Aggregate based on top-1 prediction from each classifier
        if max_aggregate:
            preds = super().predict(x, batch_size=batch_size, raw=True, **kwargs)
            aggregated_preds = np.zeros_like(preds, shape=preds.shape[1:])
            for i in range(preds.shape[0]):
                aggregated_preds[np.arange(len(aggregated_preds)), np.argmax(preds[i], axis=1)] += 1
            return aggregated_preds

        # Aggregate based on summing predictions from each classifier
        return super().predict(x, batch_size=batch_size, raw=False, **kwargs)

        Perform prediction for a batch of inputs. Aggregation will be performed on the prediction from
        each classifier if max_aggregate is True. Otherwise, the probabilities will be summed instead.
        For logits output set max_aggregate=True, as logits are not comparable between models and should
        not be aggregated using a sum.

        :param x: Input samples.
        :param batch_size: Size of batches.
        :param raw: Return the individual classifier raw outputs (not aggregated).
        :param max_aggregate: Aggregate the predicted classes of each classifier if True. If false, aggregation
               is done using a sum. If raw is true, this arg is ignored
        :return: Array of predictions of shape `(nb_inputs, nb_classes)`, or of shape
                 `(nb_classifiers, nb_inputs, nb_classes)` if `raw=True`.
        
art/estimators/classification/deep_partition_ensemble.py