Merge pull request #1971 from Trusted-AI/dependabot/pip/jax-cpu--0.4.1

Bump jax[cpu] from 0.3.25 to 0.4.1

.github/workflows/ci-tensorflow-v1.yml

@@ -45,13 +45,14 @@ jobs:
           sudo apt-get update
           sudo apt-get -y -q install ffmpeg libavcodec-extra
           python -m pip install --upgrade pip setuptools wheel
-          pip install -q -r <(sed '/^pandas/d;/^scipy/d;/^matplotlib/d;/^xgboost/d' requirements_test.txt)
+          pip install -q -r <(sed '/^pandas/d;/^scipy/d;/^matplotlib/d;/^xgboost/d;/^jax/d' requirements_test.txt)
           pip install pandas==1.3.5
           pip install scipy==1.7.2
           pip install matplotlib==3.5.3
           pip install xgboost==1.6.2
           pip install tensorflow==${{ matrix.tensorflow }}
           pip install keras==${{ matrix.keras }}
+          pip install jax[cpu]==0.3.25
           pip list
       - name: Run Tests
         run: ./run_tests.sh ${{ matrix.framework }}

requirements_test.txt

@@ -42,7 +42,7 @@ xgboost==1.7.2
 kornia~=0.6.8
 tensorboardX==2.5.1
 lief==0.12.3
-jax[cpu]==0.3.25
+jax[cpu]==0.4.1
 
 # Lingvo ASR dependencies
 # supported versions: (lingvo==0.6.4 with tensorflow-gpu==2.1.0)