def __init__(
        self,
        model: list,
        predict_func: Callable,
        loss_func: Callable,
        update_func: Callable,
        input_shape: tuple[int, ...],
        nb_classes: int,
        channels_first: bool = False,
        clip_values: "CLIP_VALUES_TYPE" | None = None,
        preprocessing_defences: "Preprocessor" | list["Preprocessor"] | None = None,
        postprocessing_defences: "Postprocessor" | list["Postprocessor"] | None = None,
        preprocessing: "PREPROCESSING_TYPE" = (0.0, 1.0),
    ) -> None:
        """
        Initialization specifically for the Jax-based implementation.

        :param model: Jax model, represented as a list of model parameters.
        :param predict_func: A function used to predict model output given the model and the input.
        :param loss_func: The loss function for which to compute gradients for training.
        :param update_func: The update function for which to train the model.
        :param input_shape: The shape of one input instance.
        :param nb_classes: The number of classes of the model.
        :param channels_first: Set channels first or last.
        :param clip_values: Tuple of the form `(min, max)` of floats or `np.ndarray` representing the minimum and
               maximum values allowed for features. If floats are provided, these will be used as the range of all
               features. If arrays are provided, each value will be considered the bound for a feature, thus
               the shape of clip values needs to match the total number of features.
        :param preprocessing_defences: Preprocessing defence(s) to be applied by the classifier.
        :param postprocessing_defences: Postprocessing defence(s) to be applied by the classifier.
        :param preprocessing: Tuple of the form `(subtrahend, divisor)` of floats or `np.ndarray` of values to be
               used for data preprocessing. The first value will be subtracted from the input. The input will then
               be divided by the second one.
        """
        super().__init__(
            model=model,
            clip_values=clip_values,
            channels_first=channels_first,
            preprocessing_defences=preprocessing_defences,
            postprocessing_defences=postprocessing_defences,
            preprocessing=preprocessing,
        )

        self._predict_func = predict_func
        self._loss_func = loss_func
        self._update_func = update_func
        self._nb_classes = nb_classes
        self._input_shape = input_shape

        Initialization specifically for the Jax-based implementation.

        :param model: Jax model, represented as a list of model parameters.
        :param predict_func: A function used to predict model output given the model and the input.
        :param loss_func: The loss function for which to compute gradients for training.
        :param update_func: The update function for which to train the model.
        :param input_shape: The shape of one input instance.
        :param nb_classes: The number of classes of the model.
        :param channels_first: Set channels first or last.
        :param clip_values: Tuple of the form `(min, max)` of floats or `np.ndarray` representing the minimum and
               maximum values allowed for features. If floats are provided, these will be used as the range of all
               features. If arrays are provided, each value will be considered the bound for a feature, thus
               the shape of clip values needs to match the total number of features.
        :param preprocessing_defences: Preprocessing defence(s) to be applied by the classifier.
        :param postprocessing_defences: Postprocessing defence(s) to be applied by the classifier.
        :param preprocessing: Tuple of the form `(subtrahend, divisor)` of floats or `np.ndarray` of values to be
               used for data preprocessing. The first value will be subtracted from the input. The input will then
               be divided by the second one.
        
art/experimental/estimators/classification/jax.py