Remove JAX constraints in Exporter (ultralytics#8309)

hmurari · Apr 17, 2024 · 969f5e1 · 969f5e1
1 parent 266940c
commit 969f5e1
Show file tree

Hide file tree

Showing 2 changed files with 1 addition and 4 deletions.
diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml
@@ -113,8 +113,6 @@ jobs:
         run: |
           python -m pip install --upgrade pip wheel
           pip install -e ".[export]" "coverage[toml]" --extra-index-url https://download.pytorch.org/whl/cpu
-          # Fix SavedModel issue "partially initialized module 'jax' has no attribute 'version' (most likely due to a circular import)" in https://github.com/google/jax/discussions/14036
-          # pip install -U 'jax!=0.4.15' 'jaxlib!=0.4.15'
           yolo export format=tflite imgsz=32 || true
       - name: Check environment
         run: |

diff --git a/ultralytics/engine/exporter.py b/ultralytics/engine/exporter.py
@@ -860,8 +860,7 @@ def export_edgetpu(self, tflite_model="", prefix=colorstr("Edge TPU:")):
     @try_export
     def export_tfjs(self, prefix=colorstr("TensorFlow.js:")):
         """YOLOv8 TensorFlow.js export."""
-        # JAX bug requiring install constraints in https://github.com/google/jax/issues/18978
-        check_requirements(["jax<=0.4.21", "jaxlib<=0.4.21", "tensorflowjs"])
+        check_requirements("tensorflowjs")
         import tensorflow as tf
         import tensorflowjs as tfjs  # noqa