diff --git a/tensorflow_gnn/runner/__init__.py b/tensorflow_gnn/runner/__init__.py
index 030d7e0f..160f3418 100644
--- a/tensorflow_gnn/runner/__init__.py
+++ b/tensorflow_gnn/runner/__init__.py
@@ -53,6 +53,7 @@
 
 # Model helpers
 chain_first_output = model_utils.chain_first_output
+integrated_gradients = attribution.integrated_gradients
 
 # Orchestration
 run = orchestration.run