Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

Remove JAX constraints in Exporter #8309

Merged
merged 2 commits into from
Feb 19, 2024
Merged

Remove JAX constraints in Exporter #8309

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
0eaedf2
Remove JAX constraints in Exporter
glenn-jocher Feb 19, 2024
7d112e6
Update ci.yaml
glenn-jocher Feb 19, 2024
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Jump to
Jump to file
Failed to load files.
Diff view
Diff view
2 changes: 0 additions & 2 deletions .github/workflows/ci.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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: |
Expand Down
3 changes: 1 addition & 2 deletions ultralytics/engine/exporter.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down