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()

        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`.
        
art/experimental/estimators/classification/jax.py