Fix jax version to 0.4.28

fineguy · The TensorFlow Datasets Authors · commit 8731daadf016 · 2024-06-11T06:55:35.000-07:00
PiperOrigin-RevId: 642252354
diff --git a/setup.py b/setup.py
@@ -86,7 +86,7 @@
 
 TESTS_DEPENDENCIES = [
     'dill',
-    'jax[cpu]',
+    'jax[cpu]==0.4.28',
     'jupyter',
     'pytest',
     'pytest-shard',
