From 9361ce6e2911bfdbf0d69028517179d4ba78fbf6 Mon Sep 17 00:00:00 2001
From: Fangjun Kuang <csukuangfj@gmail.com>
Date: Thu, 25 Apr 2024 20:53:17 +0800
Subject: [PATCH] support torch 2.3.0

---
 .github/workflows/macos-cpu-wheels.yml        | 24 +++++++++++++++----
 .../github_actions/generate_build_matrix.py   | 11 ++++++++-
 scripts/github_actions/install_torch.sh       | 13 ++++++++++
 3 files changed, 42 insertions(+), 6 deletions(-)

diff --git a/.github/workflows/macos-cpu-wheels.yml b/.github/workflows/macos-cpu-wheels.yml
index 36fb993..3f01d11 100644
--- a/.github/workflows/macos-cpu-wheels.yml
+++ b/.github/workflows/macos-cpu-wheels.yml
@@ -36,7 +36,7 @@ jobs:
   build_wheels_macos_cpu:
     needs: generate_build_matrix
     name: ${{ matrix.torch }} ${{ matrix.python-version }}
-    runs-on: macos-latest
+    runs-on: macos-14
     strategy:
       fail-fast: false
       matrix:
@@ -48,14 +48,14 @@ jobs:
           fetch-depth: 0
 
       - name: Setup Python
-        uses: actions/setup-python@v2
+        uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
 
       - name: Install dependencies
         shell: bash
         run: |
-          pip install -q torch==${{ matrix.torch}} cmake numpy wheel twine setuptools
+          pip install -q torch==${{ matrix.torch}} cmake numpy wheel>=0.40.0 twine setuptools
 
       - name: Build wheel
         shell: bash
@@ -64,13 +64,27 @@ jobs:
           mkdir wheelhouse
           cp -v dist/* wheelhouse
 
-      - name: Display wheels
+      - name: Display wheels (before fix)
+        shell: bash
+        run: |
+          ls -lh ./wheelhouse/
+
+      - name: Fix wheel platform tag
+        run: |
+          # See https://github.com/glencoesoftware/zeroc-ice-py-macos-x86_64/pull/3/files
+          # See:
+          #  * https://github.com/pypa/wheel/issues/406
+          python -m wheel tags \
+            --platform-tag=macosx_11_0_arm64 \
+            --remove wheelhouse/*.whl
+
+      - name: Display wheels (after fix)
         shell: bash
         run: |
           ls -lh ./wheelhouse/
 
       - name: Upload Wheel
-        uses: actions/upload-artifact@v2
+        uses: actions/upload-artifact@v3
         with:
           name: torch-${{ matrix.torch }}-python-${{ matrix.python-version }}-macos-latest-cpu
           path: wheelhouse/*.whl
diff --git a/scripts/github_actions/generate_build_matrix.py b/scripts/github_actions/generate_build_matrix.py
index 08ce7b9..7cb56f1 100755
--- a/scripts/github_actions/generate_build_matrix.py
+++ b/scripts/github_actions/generate_build_matrix.py
@@ -195,10 +195,16 @@ def generate_build_matrix(enable_cuda, for_windows, for_macos, test_only_latest_
             if not for_windows
             else ["11.8.0", "12.1.0"],
         },
+        "2.3.0": {
+            "python-version": ["3.8", "3.9", "3.10", "3.11", "3.12"],
+            "cuda": ["11.8", "12.1"]  # default 12.1
+            if not for_windows
+            else ["11.8.0", "12.1.0"],
+        },
         # https://github.com/Jimver/cuda-toolkit/blob/master/src/links/windows-links.ts
     }
     if test_only_latest_torch:
-        latest = "2.2.2"
+        latest = "2.3.0"
         matrix = {latest: matrix[latest]}
 
     if for_windows or for_macos:
@@ -234,6 +240,9 @@ def generate_build_matrix(enable_cuda, for_windows, for_macos, test_only_latest_
             for p in python_versions:
                 if p in excluded_python_versions:
                     continue
+                if for_macos and p in ["3.8", "3.9"]:
+                    # macOS arm64 in github actions does not support python 3.8 or 3.9
+                    continue
 
                 if for_windows:
                     p = "cp" + "".join(p.split("."))
diff --git a/scripts/github_actions/install_torch.sh b/scripts/github_actions/install_torch.sh
index d071a02..298456b 100755
--- a/scripts/github_actions/install_torch.sh
+++ b/scripts/github_actions/install_torch.sh
@@ -198,6 +198,19 @@ case ${torch} in
         ;;
     esac
     ;;
+  2.3.*)
+    case ${cuda} in
+      11.8)
+        package="torch==${torch}+cu118"
+        url=https://download.pytorch.org/whl/torch_stable.html
+        ;;
+      12.1)
+        package="torch==${torch}"
+        # Leave it empty to use PyPI.
+        url=
+        ;;
+    esac
+    ;;
   *)
     echo "Unsupported PyTorch version: ${torch}"
     exit 1