diff --git a/requirements/ai-ml_pip.txt b/requirements/ai-ml_pip.txt
index cf710ec..7f1582f 100644
--- a/requirements/ai-ml_pip.txt
+++ b/requirements/ai-ml_pip.txt
@@ -3,6 +3,7 @@ deepspeed
 wandb
 # equinox  # removed to save space
 flax
+keras
 # haliax  # removed to save space
 # levanter  # removed to save space
 # gradio # requires pydantic>=2, which conflicts with jupyter-ai. We can re-include once juptyer-ai is updated to support pydantic>=2