From 37b2c1ce18544a92a192a5d72cbff3a9532aa5e6 Mon Sep 17 00:00:00 2001
From: Chun-nien Chan <cnchan@google.com>
Date: Mon, 16 Dec 2024 10:40:01 -0800
Subject: [PATCH] enable macos unit tests

PiperOrigin-RevId: 706759813
---
 .github/workflows/nightly_unittests.yml |  1 +
 .github/workflows/unittests_python.yml  | 26 +++++++++++++++++++++----
 requirements.txt                        |  9 +++------
 3 files changed, 26 insertions(+), 10 deletions(-)

diff --git a/.github/workflows/nightly_unittests.yml b/.github/workflows/nightly_unittests.yml
index bf9907d7..44461555 100644
--- a/.github/workflows/nightly_unittests.yml
+++ b/.github/workflows/nightly_unittests.yml
@@ -20,3 +20,4 @@ jobs:
     uses: ./.github/workflows/unittests_python.yml
     with:
       trigger-sha: ${{ github.sha }}
+      run-on-macos: true
diff --git a/.github/workflows/unittests_python.yml b/.github/workflows/unittests_python.yml
index 33e6f3b9..f4a4a568 100644
--- a/.github/workflows/unittests_python.yml
+++ b/.github/workflows/unittests_python.yml
@@ -12,16 +12,34 @@ on:
   workflow_call:
     inputs:
       trigger-sha:
-        required: true
         type: string
-
+        required: true
+      run-on-macos:
+        type: boolean
+        default: true
 jobs:
+  prepare-matrix:
+    runs-on: ubuntu-latest
+    outputs:
+      os: ${{ steps.set-matrix.outputs.os }}
+    steps:
+      - name: Set matrix
+        id: set-matrix
+        run: |
+          os=("Linux_runner_8_core")
+          if [[ "${{ inputs.run-on-macos }}" == "true" ]]; then
+            os+=("macos-latest")
+          fi
+          os=$(jq -c -n '$ARGS.positional' --args "${os[@]}")
+          echo "os=${os}"
+          echo "os=${os}" >> $GITHUB_OUTPUT
   test:
+    needs: prepare-matrix
     strategy:
       matrix:
         python-version: ["3.10", "3.11"]
-    runs-on:
-      labels: Linux_runner_8_core
+        os: ${{ fromJson(needs.prepare-matrix.outputs.os) }}
+    runs-on: ${{ matrix.os }}
     steps:
       - uses: actions/checkout@v3
         with:
diff --git a/requirements.txt b/requirements.txt
index 129b0ebe..c94c6106 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,9 +1,6 @@
--f https://download.pytorch.org/whl/torch/
-torch==2.5.1+cpu
--f https://download.pytorch.org/whl/torchvision/
-torchvision==0.20.1+cpu
--f https://download.pytorch.org/whl/torchaudio/
-torchaudio==2.5.1+cpu
+torch==2.5.1
+torchvision==0.20.1
+torchaudio==2.5.1
 --pre
 tf-nightly>=2.19.0.dev20241201
 ai-edge-litert-nightly