Need include jax in requirements.txt after all (#994)

The requirements of tensorflow and tf-keras _should_ bring in a
compatible jax version; however, in some environments, it doesn't happen
as expected, and `pip install -r requirements.txt` results in a
confusing error about trying to use NumPy 2 with a package built for
NumPy 1.

Author: mhucka

requirements.in

@@ -19,6 +19,7 @@
 cirq-core<=1.5.0
 cirq-google<=1.5.0
 contourpy<=1.3.2
+jax>=0.5,<0.6
 numpy<=2.2.6
 scipy<1.16
 sympy==1.14

requirements.txt

@@ -82,6 +82,10 @@ ipython-genutils==0.2.0
     # via nbformat
 isort==5.13.2
     # via pylint
+jax==0.5.3
+    # via -r requirements.in
+jaxlib==0.5.3
+    # via jax
 jinja2==3.1.6
     # via tensorflow-docs
 jsonschema==4.26.0
@@ -116,6 +120,8 @@ mdurl==0.1.2
     # via markdown-it-py
 ml-dtypes==0.5.4
     # via
+    #   jax
+    #   jaxlib
     #   keras
     #   tensorflow
 mpmath==1.3.0
@@ -139,6 +145,8 @@ numpy==2.0.2
     #   cirq-core
     #   contourpy
     #   h5py
+    #   jax
+    #   jaxlib
     #   keras
     #   matplotlib
     #   ml-dtypes
@@ -148,7 +156,9 @@ numpy==2.0.2
     #   tensorflow
     #   typedunits
 opt-einsum==3.4.0
-    # via tensorflow
+    # via
+    #   jax
+    #   tensorflow
 optree==0.18.0
     # via keras
 packaging==26.0
@@ -227,6 +237,8 @@ scipy==1.15.3
     # via
     #   -r requirements.in
     #   cirq-core
+    #   jax
+    #   jaxlib
 setuptools==75.0.0
     # via
     #   tensorboard