Unpin JAX and Tensorflow versions.

To get the latest ones: JAX 0.5.2 and Tensorflow 2.19.0

Author: hertschuh

integration_tests/import_test.py

@@ -28,11 +28,11 @@ def setup_package():
     whl_path = re.findall(
         r"[^\s]*\.whl",
         build_process.stdout,
-    )[-1]
+    )
     if not whl_path:
         print(build_process.stderr)
         raise ValueError("Installing Keras package unsuccessful. ")
-    return whl_path
+    return whl_path[-1]
 
 
 def create_virtualenv():

keras/src/backend/jax/image.py

@@ -465,7 +465,7 @@ def affine_transform(
     # transform the indices
     coordinates = jnp.einsum("Bhwij, Bjk -> Bhwik", indices, transform)
     coordinates = jnp.moveaxis(coordinates, source=-1, destination=1)
-    coordinates += jnp.reshape(a=offset, shape=(*offset.shape, 1, 1, 1))
+    coordinates += jnp.reshape(offset, shape=(*offset.shape, 1, 1, 1))
 
     # apply affine transformation
     _map_coordinates = functools.partial(

keras/src/backend/torch/image.py

@@ -424,7 +424,7 @@ def affine_transform(
     # transform the indices
     coordinates = torch.einsum("Bhwij, Bjk -> Bhwik", indices, transform)
     coordinates = torch.moveaxis(coordinates, source=-1, destination=1)
-    coordinates += torch.reshape(a=offset, shape=(*offset.shape, 1, 1, 1))
+    coordinates += torch.reshape(offset, shape=(*offset.shape, 1, 1, 1))
 
     # Note: torch.stack is faster than torch.vmap when the batch size is small.
     affined = torch.stack(

keras/src/export/onnx.py

@@ -80,7 +80,7 @@ def export_onnx(model, filepath, verbose=None, input_signature=None, **kwargs):
         decorated_fn = get_concrete_fn(model, input_signature, **kwargs)
 
         # Use `tf2onnx` to convert the `decorated_fn` to the ONNX format.
-        patch_tf2onnx()  # TODO: Remove this once `tf2onnx` supports numpy 2.
+        # patch_tf2onnx()  # TODO: Remove this once `tf2onnx` supports numpy 2.
         tf2onnx.convert.from_function(
             decorated_fn, input_signature, output_path=filepath
         )

requirements-common.txt

@@ -10,13 +10,14 @@ absl-py
 requests
 h5py
 ml-dtypes
-protobuf
+protobuf>=4.21.6  # Earlier versions break Tensorflow>=2.19
 tensorboard-plugin-profile
 rich
 build
 optree
 pytest-cov
 packaging
+tf2onnx>=1.16.1  # For Numpy 2 support
 # for tree_test.py
 dm_tree
 coverage!=7.6.5  # 7.6.5 breaks CI

requirements-jax-cuda.txt

@@ -1,14 +1,13 @@
 # Tensorflow cpu-only version (needed for testing).
-tensorflow-cpu~=2.18.0
-tf2onnx
+tensorflow-cpu
 
 # Torch cpu-only version (needed for testing).
 --extra-index-url https://download.pytorch.org/whl/cpu
 torch==2.6.0+cpu
 
 # Jax with cuda support.
 --find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
-jax[cuda12]==0.4.28
+jax[cuda12]
 flax
 
 -r requirements-common.txt

requirements-tensorflow-cuda.txt

@@ -1,6 +1,5 @@
 # Tensorflow with cuda support.
-tensorflow[and-cuda]~=2.18.0
-tf2onnx
+tensorflow[and-cuda]
 
 # Torch cpu-only version (needed for testing).
 --extra-index-url https://download.pytorch.org/whl/cpu

requirements-torch-cuda.txt

@@ -1,6 +1,5 @@
 # Tensorflow cpu-only version (needed for testing).
 tensorflow-cpu~=2.18.0
-tf2onnx
 
 # Torch with cuda support.
 # - torch is pinned to a version that is compatible with torch-xla

requirements.txt

@@ -1,18 +1,15 @@
 # Tensorflow.
-tensorflow-cpu~=2.18.0;sys_platform != 'darwin'
-tensorflow~=2.18.0;sys_platform == 'darwin'
+tensorflow-cpu~=2.19.0;sys_platform != 'darwin'
+tensorflow~=2.19.0;sys_platform == 'darwin'
 tf_keras
-tf2onnx
 
 # Torch.
 --extra-index-url https://download.pytorch.org/whl/cpu
 torch==2.6.0+cpu
 torch-xla==2.6.0;sys_platform != 'darwin'
 
 # Jax.
-# Pinned to 0.5.0 on CPU. JAX 0.5.1 requires Tensorflow 2.19 for saved_model_test.
-# Note that we test against the latest JAX on GPU.
-jax[cpu]==0.5.0
+jax[cpu]
 flax
 
 # Common deps.