Pin `jax` and `jaxlib` to fix `tensorflowjs` exports (#6988)

Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Glenn Jocher <glenn.jocher@ultralytics.com>

requirements.txt

@@ -32,6 +32,8 @@ seaborn>=0.11.0
 # scikit-learn==0.19.2  # CoreML quantization
 # tensorflow>=2.4.1,<=2.13.1  # TF exports (-cpu, -aarch64, -macos)
 # tflite-support
+# jax<=0.4.21  # tensorflowjs bug https://github.com/google/jax/issues/18978
+# jaxlib<=0.4.21  # tensorflowjs bug https://github.com/google/jax/issues/18978
 # tensorflowjs>=3.9.0  # TF.js export
 # openvino-dev>=2023.0  # OpenVINO export
 

setup.py

@@ -81,6 +81,8 @@ setup(
             'coremltools>=7.0',
             'openvino-dev>=2023.0',
             'tensorflow<=2.13.1',  # TF bug https://github.com/ultralytics/ultralytics/issues/5161
+            'jax<=0.4.21',  # tensorflowjs bug https://github.com/google/jax/issues/18978
+            'jaxlib<=0.4.21',  # tensorflowjs bug https://github.com/google/jax/issues/18978
             'tensorflowjs',  # automatically installs tensorflow
         ], },
     classifiers=[

ultralytics/engine/exporter.py

@@ -781,7 +781,8 @@ class Exporter:
     @try_export
     def export_tfjs(self, prefix=colorstr('TensorFlow.js:')):
         """YOLOv8 TensorFlow.js export."""
-        check_requirements('tensorflowjs')
+        # JAX bug requiring install constraints in https://github.com/google/jax/issues/18978
+        check_requirements(['jax<=0.4.21', 'jaxlib<=0.4.21', 'tensorflowjs'])
         import tensorflow as tf
         import tensorflowjs as tfjs  # noqa