From e4a62f99f767cfe0f5279ffae1a1f2eaf63af9ba Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 19 Feb 2024 19:44:51 +0100 Subject: [PATCH] Remove JAX constraints in Exporter (#8309) --- .github/workflows/ci.yaml | 2 -- ultralytics/engine/exporter.py | 3 +-- 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 60c0486e..e8019f4d 100644 --- 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 index b5900840..ff3ef05f 100644 --- a/ultralytics/engine/exporter.py +++ b/ultralytics/engine/exporter.py @@ -860,8 +860,7 @@ class Exporter: @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