diff --git a/docker/requirement_mappings.json b/docker/requirement_mappings.json
index 09d17a4874299..e87d8e8b671e6 100644
--- a/docker/requirement_mappings.json
+++ b/docker/requirement_mappings.json
@@ -1,6 +1,6 @@
 {
-  "tensorflow": ["tensorflow-cpu==2.15.1", "tensorflow-probability"],
-  "jax": ["ml-dtypes","jax[cpu]","dm-haiku", "flax", "jaxlib"],
+  "tensorflow": ["tensorflow", "tensorflow-probability"],
+  "jax": ["ml-dtypes", "jax[cpu]", "dm-haiku", "flax", "jaxlib"],
   "numpy": ["numpy"],
   "paddle": ["paddlepaddle"],
   "mxnet": ["mxnet"],