Bugfix in concatenating gradient batches

tf.stack only accounted for the last batch of gradients, instead of collecting the gradients from all batches. The all but the last batch were discarded. The correct implementation is to concatenate the gradient batches by the first axis. This way, all the gradients are correctly accounted for.

Tested the changes by running it in googe colab.

Author: mhavasi

site/en/tutorials/interpretability/integrated_gradients.ipynb

@@ -904,8 +904,8 @@
         "    gradient_batch = one_batch(baseline, image, alpha_batch, target_class_idx)\n",
         "    gradient_batches.append(gradient_batch)\n",
         "      \n",
-        "  # Stack path gradients together row-wise into single tensor.\n",
-        "  total_gradients = tf.stack(gradient_batch)\n",
+        "  # Concatenate path gradients together row-wise into single tensor.\n",
+        "  total_gradients = tf.concat(gradient_batches, axis=0)\n",
         "\n",
         "  # Integral approximation through averaging gradients.\n",
         "  avg_gradients = integral_approximation(gradients=total_gradients)\n",