Update TensorFlow ml_dtypes dependency to >= 0.3.1 < 0.5.0

PiperOrigin-RevId: 642337465

ci/official/requirements_updater/requirements.in

@@ -13,6 +13,7 @@ gast == 0.4.0
 termcolor == 2.3.0
 wrapt == 1.16.0
 tblib == 2.0.0
+ml_dtypes >= 0.4.0, < 0.5.0
 
 # Install tensorboard, and keras
 # Note that here we want the latest version that matches TF major.minor version

requirements_lock_3_10.txt

@@ -346,6 +346,7 @@ ml-dtypes==0.4.0 \
     --hash=sha256:ee9f91d4c4f9959a7e1051c141dc565f39e54435618152219769e24f5e9a4d06 \
     --hash=sha256:f1724ddcdf5edbaf615a62110af47407f1719b8d02e68ccee60683acb5f74da1
     # via
+    #   -r requirements.in
     #   jax
     #   keras-nightly
 namex==0.0.8 \

requirements_lock_3_11.txt

@@ -346,6 +346,7 @@ ml-dtypes==0.4.0 \
     --hash=sha256:ee9f91d4c4f9959a7e1051c141dc565f39e54435618152219769e24f5e9a4d06 \
     --hash=sha256:f1724ddcdf5edbaf615a62110af47407f1719b8d02e68ccee60683acb5f74da1
     # via
+    #   -r requirements.in
     #   jax
     #   keras-nightly
 namex==0.0.8 \

requirements_lock_3_12.txt

@@ -346,6 +346,7 @@ ml-dtypes==0.4.0 \
     --hash=sha256:ee9f91d4c4f9959a7e1051c141dc565f39e54435618152219769e24f5e9a4d06 \
     --hash=sha256:f1724ddcdf5edbaf615a62110af47407f1719b8d02e68ccee60683acb5f74da1
     # via
+    #   -r requirements.in
     #   jax
     #   keras-nightly
 namex==0.0.8 \

requirements_lock_3_9.txt

@@ -350,6 +350,7 @@ ml-dtypes==0.4.0 \
     --hash=sha256:ee9f91d4c4f9959a7e1051c141dc565f39e54435618152219769e24f5e9a4d06 \
     --hash=sha256:f1724ddcdf5edbaf615a62110af47407f1719b8d02e68ccee60683acb5f74da1
     # via
+    #   -r requirements.in
     #   jax
     #   keras-nightly
 namex==0.0.8 \

tensorflow/tools/pip_package/setup.py

@@ -83,7 +83,7 @@ REQUIRED_PACKAGES = [
     'google_pasta >= 0.1.1',
     'h5py >= 3.10.0',
     'libclang >= 13.0.0',
-    'ml_dtypes ~= 0.3.1',
+    'ml_dtypes >= 0.3.1, < 0.5.0',
     # TODO(b/304751256): Adjust the numpy pin to a single version, when ready
     'numpy >= 1.23.5, < 2.0.0 ; python_version <= "3.11"',
     'numpy >= 1.26.0, < 2.0.0 ; python_version >= "3.12"',