diff --git a/tools/testing/source_code_test.py b/tools/testing/source_code_test.py
index 5a546a9f48..a6d140fbdd 100644
--- a/tools/testing/source_code_test.py
+++ b/tools/testing/source_code_test.py
@@ -18,6 +18,7 @@
 
 from typedapi import ensure_api_is_typed
 
+import importlib
 import tensorflow_addons as tfa
 import tensorflow as tf
 
@@ -40,8 +41,10 @@ def test_api_typed():
     # Files within this list will be exempt from verification.
     exception_list = [
         tfa.rnn.PeepholeLSTMCell,
-        tf.keras.optimizers.legacy.optimizer.Optimizer,
     ]
+    if importlib.util.find_spec("tensorflow.keras.optimizers.legacy") is not None:
+        exception_list.append(tf.keras.optimizers.legacy.optimizer.Optimizer)
+
     help_message = (
         "You can also take a look at the section about it in the CONTRIBUTING.md:\n"
         "https://github.com/tensorflow/addons/blob/master/CONTRIBUTING.md#about-type-hints"