class_docstr: This class implements a classifier with the Jax framework.
class_docstr_tok: ["This","class","implements","a","classifier","with","the","Jax","framework","."]
docstr: Compute the gradient of the loss function w.r.t. `x`.
  
  :param x: Sample input with shape as expected by the model.
  :param y: Target values.
  :return: Array of gradients of the same shape as `x`.
docstr_tok: ["Compute","the","gradient","of","the","loss","function","w.r.t",".","`","x","`",".",":","param","x",":","Sample","input","with","shape","as","expected","by","the","model",".",":","param","y",":","Target","values",".",":","return",":","Array","of","gradients","of","the","same","shape","as","`","x","`","."]
code: def loss_gradient(
      self, x: np.ndarray, y: np.ndarray, **kwargs
  ) -> np.ndarray:
      """
      Compute the gradient of the loss function w.r.t. `x`.
  
      :param x: Sample input with shape as expected by the model.
      :param y: Target values.
      :return: Array of gradients of the same shape as `x`.
      """
      from jax import grad
  
      # Apply preprocessing
      x_preprocessed, y_preprocessed = self._apply_preprocessing(
          x, y, fit=False
      )
  
      # Compute gradients
      grads = grad(self.loss_func, argnums=(0, 1))(
          self.model, x_preprocessed, y_preprocessed
      )[1]
  
      assert grads.shape == x.shape
  
      return grads.copy()
  
imports: from __future__ import absolute_import, division, print_function, unicode_literals
  import logging
  import random
  from typing import Any, Callable, Dict, List, Optional, Tuple, Union, TYPE_CHECKING
  import numpy
  from art.estimators.classification.classifier import ClassGradientsMixin, ClassifierMixin
  from art.experimental.estimators.jax import JaxEstimator
label_desc: Use image node_id 9 for calling the JaxClassifier obj's underlying member method code with example usage: obj.loss_gradient(x, y) and returns: grads