class_docstr: Estimator class for Jax models.
docstr: Perform prediction of the neural network for samples `x`.
  
  :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 batch_size: Batch size.
  :return: Predictions.
  :rtype: Format as expected by the `model`
code: def predict(self, x: np.ndarray, batch_size: int = 128, **kwargs):
      """
      Perform prediction of the neural network for samples `x`.
  
      :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 batch_size: Batch size.
      :return: Predictions.
      :rtype: Format as expected by the `model`
      """
      return NeuralNetworkMixin.predict(
          self, x, batch_size=batch_size, **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 2 for calling the JaxEstimator obj's underlying member method code with example usage: obj.predict(x, batch_size) and returns: NeuralNetworkMixin