class_docstr: Estimator class for Jax models.
docstr: Fit the model of the estimator on the training data `x` and `y`.
  
  :param x: Samples of shape (nb_samples, nb_features) or (nb_samples, nb_pixels_1, nb_pixels_2,
            nb_channels) or (nb_samples, nb_channels, nb_pixels_1, nb_pixels_2).
  :param y: Target values.
  :type y: Format as expected by the `model`
  :param batch_size: Batch size.
  :param nb_epochs: Number of training epochs.
code: def fit(
      self,
      x: np.ndarray,
      y,
      batch_size: int = 128,
      nb_epochs: int = 20,
      **kwargs,
  ) -> None:
      """
      Fit the model of the estimator on the training data `x` and `y`.
  
      :param x: Samples of shape (nb_samples, nb_features) or (nb_samples, nb_pixels_1, nb_pixels_2,
                nb_channels) or (nb_samples, nb_channels, nb_pixels_1, nb_pixels_2).
      :param y: Target values.
      :type y: Format as expected by the `model`
      :param batch_size: Batch size.
      :param nb_epochs: Number of training epochs.
      """
      NeuralNetworkMixin.fit(
          self,
          x,
          y,
          batch_size=batch_size,
          nb_epochs=nb_epochs,
          **kwargs,
      )
  
file_path: adversarial-robustness-toolbox/art/experimental/estimators/jax.py
imports: import logging
  import numpy
  from art.estimators.estimator import BaseEstimator, LossGradientsMixin, NeuralNetworkMixin
label_desc: Use image node_id 3 for calling the JaxEstimator obj's underlying member method code with example usage: obj.fit(x, y, batch_size, nb_epochs) without return types