diff --git a/pyproject.toml b/pyproject.toml
index 831d06e8..d506f9d8 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -40,7 +40,6 @@ dependencies = [
   "six>=1.16.0",
   "termcolor>=2.3.0",
   "tqdm>=4.64.0",
-  "umap-learn>=0.5.1",
   "werkzeug>=2.2.3",
 ]
 # LINT.ThenChange(./requirements.txt)
@@ -105,10 +104,14 @@ examples-generative-ai = [
   "vertexai>=1.49.0",
 ]
 # LINT.ThenChange(./requirements_examples_generative_ai.txt)
+umap = [
+  "umap-learn==0.5.6"
+]
 # LINT.IfChange
 test = [
   "lit-nlp[examples-discriminative-ai]",
   "lit-nlp[examples-generative-ai]",
+  "lit-nlp[umap]",
   "lime==0.2.0.1",
   "pytest>=7.4.0,<8.0.0",
   "webtest>=2.0",
diff --git a/requirements.txt b/requirements.txt
index 6e740ce9..9c4707b0 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -37,6 +37,5 @@ shap>=0.42.0,<0.46.0
 six>=1.16.0
 termcolor>=2.3.0
 tqdm>=4.64.0
-umap-learn>=0.5.1
 werkzeug>=2.2.3
 # LINT.ThenChange(./pyproject.toml)
diff --git a/requirements_test.txt b/requirements_test.txt
index 367a238f..9e2e0556 100644
--- a/requirements_test.txt
+++ b/requirements_test.txt
@@ -18,5 +18,6 @@
 
 lime==0.2.0.1
 pytest>=7.4.0,<8.0.0
+umap-learn==0.5.6
 webtest>=2.0
 # LINT.ThenChange(./pyproject.toml)