code: def expected_gradients_multi_output(
      inputs, references, model, num_classes, k=100
  ):
      """
      Given a batch of inputs and labels, and a model,
      symbolically computes expected gradients with k references. Unlike
      expected_gradients_full, this function is used when you want the
      expected gradients values with respect to all output classes, not just a single one.
  
      Args:
          inputs: A [batch_size, ...]-shaped tensor. The input to a model.
          references: A numpy array representing background training data to sample from.
          model:  A tf.keras.Model object, or a subclass object thereof.
          num_classes: The number of classes to take expected gradients with respect to.
          k: The number of samples to use when computing expected gradients.
      Returns:
          A tensor the same shape as the input representing the expected gradients
          feature attributions with respect to the output predictions.
      """
      eg_array = []
  
      for i in range(tf.shape(inputs)[0]):
          sample_indices = np.random.choice(
              references.shape[0], size=k, replace=False
          )
          sample_references = references[sample_indices]
  
          predictions_indexed = []
          alphas = tf.random.uniform(
              shape=(k, 1, 1, 1),
              minval=0.0,
              maxval=1.0,
              dtype=tf.float32,
          )
          current_input = tf.expand_dims(inputs[i], axis=0)
          interpolated_inputs = (
              alphas * current_input
              + (1.0 - alphas) * sample_references
          )
  
          with tf.GradientTape(persistent=True) as tape:
              tape.watch(interpolated_inputs)
  
              predictions = model(interpolated_inputs, training=False)
              for output_class in range(num_classes):
                  predictions_indexed.append(
                      predictions[:, output_class]
                  )
  
          sample_eg_array = []
          for output_class in range(num_classes):
              input_gradients = tape.gradient(
                  predictions_indexed[output_class], interpolated_inputs
              )
              difference_from_reference = (
                  current_input - sample_references
              )
              expected_gradients_samples = (
                  input_gradients * difference_from_reference
              )
              expected_gradients = tf.reduce_mean(
                  expected_gradients_samples, axis=0
              )
              sample_eg_array.append(expected_gradients)
  
          del tape
          eg_array.append(tf.stack(sample_eg_array, axis=0))
  
      return tf.stack(eg_array, axis=0)
  
file_path: attributionpriors/attributionpriors/eager_ops.py
imports: import tensorflow
  import numpy
label_desc: Use image node_id 6 for calling a global function with example usage: expected_gradients_multi_output(inputs, references, model, num_classes, k) and returns: tf