def class_gradient(
        self,
        x: np.ndarray,
        label: int | list[int] | np.ndarray | None = None,
        training_mode: bool = False,
        **kwargs,
    ) -> np.ndarray:
        """
        Compute per-class derivatives w.r.t. `x`.

        :param x: Sample input with shape as expected by the model.
        :param label: Index of a specific per-class derivative. If an integer is provided, the gradient of that class
                      output is computed for all samples. If multiple values as provided, the first dimension should
                      match the batch size of `x`, and each value will be used as target for its corresponding sample in
                      `x`. If `None`, then gradients for all classes will be computed for each sample.
        :param training_mode: `True` for model set to training mode and `False` for model set to evaluation mode.
        :return: Array of gradients of input features w.r.t. each class in the form
                 `(batch_size, nb_classes, input_shape)` when computing for all classes, otherwise shape becomes
                 `(batch_size, 1, input_shape)` when `label` parameter is specified.
        """
        import tensorflow as tf

        x = tf.convert_to_tensor(x)

        with tf.GradientTape(persistent=True) as tape:
            # Apply preprocessing
            if self.all_framework_preprocessing:
                x_grad = tf.convert_to_tensor(x)
                tape.watch(x_grad)
                x_input, _ = self._apply_preprocessing(x_grad, y=None, fit=False)
            else:
                x_preprocessed, _ = self._apply_preprocessing(x, y=None, fit=False)
                x_grad = tf.convert_to_tensor(x_preprocessed)
                tape.watch(x_grad)
                x_input = x_grad

            tape.watch(x_input)

            # Compute the gradients
            if tf.executing_eagerly():
                if label is None:
                    # Compute the gradients w.r.t. all classes
                    class_gradients = []

                    for i in range(self.nb_classes):
                        predictions = self.model(x_input, training=training_mode)
                        prediction = predictions[:, i]
                        tape.watch(prediction)

                        class_gradient = tape.gradient(prediction, x_input).numpy()
                        class_gradients.append(class_gradient)
                        # Break after 1 iteration for binary classification case
                        if len(predictions.shape) == 1 or predictions.shape[1] == 1:
                            break

                    gradients = np.swapaxes(np.array(class_gradients), 0, 1)

                elif isinstance(label, int):
                    # Compute the gradients only w.r.t. the provided label
                    predictions = self.model(x_input, training=training_mode)
                    prediction = predictions[:, label]
                    tape.watch(prediction)

                    class_gradient = tape.gradient(prediction, x_grad).numpy()
                    gradients = np.expand_dims(class_gradient, axis=1)

                else:
                    # For each sample, compute the gradients w.r.t. the indicated target class (possibly distinct)
                    class_gradients = []
                    unique_labels = list(np.unique(label))

                    for unique_label in unique_labels:
                        predictions = self.model(x_input, training=training_mode)
                        prediction = predictions[:, unique_label]
                        tape.watch(prediction)

                        class_gradient = tape.gradient(prediction, x_grad).numpy()
                        class_gradients.append(class_gradient)

                    gradients = np.swapaxes(np.array(class_gradients), 0, 1)
                    lst = [unique_labels.index(i) for i in label]
                    gradients = np.expand_dims(gradients[np.arange(len(gradients)), lst], axis=1)

                if not self.all_framework_preprocessing:
                    gradients = self._apply_preprocessing_gradient(x, gradients)

            else:
                raise NotImplementedError("Expecting eager execution.")

        return gradients

        Compute per-class derivatives w.r.t. `x`.

        :param x: Sample input with shape as expected by the model.
        :param label: Index of a specific per-class derivative. If an integer is provided, the gradient of that class
                      output is computed for all samples. If multiple values as provided, the first dimension should
                      match the batch size of `x`, and each value will be used as target for its corresponding sample in
                      `x`. If `None`, then gradients for all classes will be computed for each sample.
        :param training_mode: `True` for model set to training mode and `False` for model set to evaluation mode.
        :return: Array of gradients of input features w.r.t. each class in the form
                 `(batch_size, nb_classes, input_shape)` when computing for all classes, otherwise shape becomes
                 `(batch_size, 1, input_shape)` when `label` parameter is specified.
        
art/estimators/classification/tensorflow.py