Split Jax training dependencies (especially Scenic and TF-IO) into a …

…separate dependency group.

PiperOrigin-RevId: 614682564

.github/workflows/ci.yml

@@ -29,6 +29,6 @@ jobs:
         run: |
           sudo apt-get update
           sudo apt-get install libsndfile1 ffmpeg
-          poetry install
+          poetry install --with jaxtrain
       - name: Test with unittest
         run: poetry run python -m unittest discover -s chirp/tests -p "*test.py"

.github/workflows/ci_no_jaxtrain.yml

@@ -0,0 +1,35 @@
+name: CI_no_jaxtrain
+
+on:
+  push:
+    branches: [ main ]
+  pull_request:
+    branches: [ main ]
+  # Allows to run this workflow manually from the Actions tab on GitHub.
+  workflow_dispatch:
+
+jobs:
+  test-ubuntu:
+    name: "test on ${{ matrix.python-version }} on ${{ matrix.os }}"
+    runs-on: "${{ matrix.os }}"
+    strategy:
+      matrix:
+        python-version: ["3.10", "3.11"]
+        os: [ubuntu-latest]
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Poetry
+        run: pipx install poetry
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+          cache: 'poetry'
+      - name: Install Chirp and its dependencies, except jaxtrain dependencies.
+        run: |
+          sudo apt-get update
+          sudo apt-get install libsndfile1 ffmpeg
+          poetry install --without jaxtrain
+      - name: Test with unittest
+        # TODO: Group together jaxtrain tests so they can be easily excluded.
+        run: poetry run python -m unittest discover -s chirp/tests -p "*inference_test.py"

chirp/inference/colab_utils.py

@@ -19,7 +19,6 @@
 
 from absl import logging
 from chirp import config_utils
-from chirp.configs import config_globals
 from chirp.inference import embed_lib
 import numpy as np
 import tensorflow as tf