From 2506f7b25742f2734a782aaa3a371bc80f2868d6 Mon Sep 17 00:00:00 2001
From: Dhruv Nair <dhruv.nair@gmail.com>
Date: Tue, 7 Nov 2023 06:03:48 +0000
Subject: [PATCH 1/4] explicit torch dependency check

---
 .../workflows/pr_torch_dependency_test.yml    | 32 +++++++++++++++++++
 1 file changed, 32 insertions(+)
 create mode 100644 .github/workflows/pr_torch_dependency_test.yml

diff --git a/.github/workflows/pr_torch_dependency_test.yml b/.github/workflows/pr_torch_dependency_test.yml
new file mode 100644
index 000000000000..3a4bb2811187
--- /dev/null
+++ b/.github/workflows/pr_torch_dependency_test.yml
@@ -0,0 +1,32 @@
+name: Run Torch dependency tests
+
+on:
+  pull_request:
+    branches:
+      - main
+  push:
+    branches:
+      - main
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+
+jobs:
+  check_dependencies:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.8"
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -e .
+          pip install torch torchvision
+          pip install pytest
+      - name: Check for soft dependencies
+        run: |
+          pytest tests/others/test_dependencies.py

From 0449ee1c946c7d61f60f5eaa229d8d1460e4d846 Mon Sep 17 00:00:00 2001
From: Dhruv Nair <dhruv.nair@gmail.com>
Date: Tue, 7 Nov 2023 06:47:00 +0000
Subject: [PATCH 2/4] update

---
 .github/workflows/pr_flax_dependency_test.yml | 35 +++++++++++++++++++
 .../workflows/pr_torch_dependency_test.yml    |  4 +--
 2 files changed, 37 insertions(+), 2 deletions(-)
 create mode 100644 .github/workflows/pr_flax_dependency_test.yml

diff --git a/.github/workflows/pr_flax_dependency_test.yml b/.github/workflows/pr_flax_dependency_test.yml
new file mode 100644
index 000000000000..60c4a5098958
--- /dev/null
+++ b/.github/workflows/pr_flax_dependency_test.yml
@@ -0,0 +1,35 @@
+name: Run Flax dependency tests
+
+on:
+  pull_request:
+    branches:
+      - main
+  push:
+    branches:
+      - main
+
+concurrency:
+  group: ${{ github.workflow }}-${{ github.head_ref || github.run_id }}
+  cancel-in-progress: true
+
+jobs:
+  check_flax_dependencies:
+    runs-on: diffusers-pytorch-dependency-cpu
+    steps:
+      - uses: actions/checkout@v3
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        with:
+          python-version: "3.8"
+      - name: Install dependencies
+        run: |
+          python -m pip install --upgrade pip
+          pip install -e .
+          pip install "jax[cpu]>=0.2.16,!=0.3.2"
+          pip install --upgrade clu
+          git clone https://github.com/google/flax.git
+          pip install --user -e flax
+          pip install pytest
+      - name: Check for soft dependencies
+        run: |
+          pytest tests/others/test_dependencies.py
diff --git a/.github/workflows/pr_torch_dependency_test.yml b/.github/workflows/pr_torch_dependency_test.yml
index 3a4bb2811187..57a7a5c77c74 100644
--- a/.github/workflows/pr_torch_dependency_test.yml
+++ b/.github/workflows/pr_torch_dependency_test.yml
@@ -13,7 +13,7 @@ concurrency:
   cancel-in-progress: true
 
 jobs:
-  check_dependencies:
+  check_torch_dependencies:
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v3
@@ -25,7 +25,7 @@ jobs:
         run: |
           python -m pip install --upgrade pip
           pip install -e .
-          pip install torch torchvision
+          pip install torch torchvision torchaudio
           pip install pytest
       - name: Check for soft dependencies
         run: |

From 1433ce3bc4937f8639d292cc63fecb4c603ea123 Mon Sep 17 00:00:00 2001
From: Dhruv Nair <dhruv.nair@gmail.com>
Date: Tue, 7 Nov 2023 06:48:21 +0000
Subject: [PATCH 3/4] update

---
 .github/workflows/pr_flax_dependency_test.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/pr_flax_dependency_test.yml b/.github/workflows/pr_flax_dependency_test.yml
index 60c4a5098958..9ed70f934692 100644
--- a/.github/workflows/pr_flax_dependency_test.yml
+++ b/.github/workflows/pr_flax_dependency_test.yml
@@ -14,7 +14,7 @@ concurrency:
 
 jobs:
   check_flax_dependencies:
-    runs-on: diffusers-pytorch-dependency-cpu
+    runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v3
       - name: Set up Python

From e3d8de88e78c1d76ce91756f697012cc9290f423 Mon Sep 17 00:00:00 2001
From: Dhruv Nair <dhruv.nair@gmail.com>
Date: Tue, 7 Nov 2023 07:05:52 +0000
Subject: [PATCH 4/4] update

---
 .github/workflows/pr_flax_dependency_test.yml | 5 ++---
 1 file changed, 2 insertions(+), 3 deletions(-)

diff --git a/.github/workflows/pr_flax_dependency_test.yml b/.github/workflows/pr_flax_dependency_test.yml
index 9ed70f934692..d7d2a2d4c3d5 100644
--- a/.github/workflows/pr_flax_dependency_test.yml
+++ b/.github/workflows/pr_flax_dependency_test.yml
@@ -26,9 +26,8 @@ jobs:
           python -m pip install --upgrade pip
           pip install -e .
           pip install "jax[cpu]>=0.2.16,!=0.3.2"
-          pip install --upgrade clu
-          git clone https://github.com/google/flax.git
-          pip install --user -e flax
+          pip install "flax>=0.4.1"
+          pip install "jaxlib>=0.1.65"
           pip install pytest
       - name: Check for soft dependencies
         run: |