Add upload wheels file for pypi (#184)

build/rocm/upload_wheels.sh

@@ -0,0 +1,53 @@
+#!/bin/bash
+
+# Check for user-supplied arguments.
+if [[ $# -lt 2 ]]; then
+    echo "Usage: $0 <jax_home_directory> <version>"
+    exit 1
+fi
+
+# Set JAX_HOME and RELEASE_VERSION from user arguments.
+JAX_HOME=$1
+RELEASE_VERSION=$2
+WHEELHOUSE="$JAX_HOME/wheelhouse"
+
+# Projects to upload separately to PyPI.
+PROJECTS=("jax_rocm60_pjrt" "jax_rocm60_plugin")
+
+# PyPI API Token.
+PYPI_API_TOKEN=${PYPI_API_TOKEN:-"pypi-replace_with_token"}
+
+# Ensure the specified JAX_HOME and wheelhouse directories exists.
+if [[ ! -d "$JAX_HOME" ]]; then
+    echo "Error: The specified JAX_HOME directory does not exist: $JAX_HOME"
+    exit 1
+fi
+if [[ ! -d "$WHEELHOUSE" ]]; then
+    echo "Error: The wheelhouse directory does not exist: $WHEELHOUSE"
+    exit 1
+fi
+
+upload_and_release_project() {
+    local project=$1
+
+    echo "Searching for wheels matching project: $project version: $RELEASE_VERSION..."
+    wheels=($(ls $WHEELHOUSE | grep "^${project}-${RELEASE_VERSION}[.-].*\.whl"))
+    if [[ ${#wheels[@]} -eq 0 ]]; then
+        echo "No wheels found for project: $project version: $RELEASE_VERSION. Skipping..."
+        return
+    fi
+    echo "Found wheels for $project: ${wheels[*]}"
+
+    echo "Uploading wheels for $project version $RELEASE_VERSION to PyPI..."
+    for wheel in "${wheels[@]}"; do
+        twine upload --verbose --repository pypi --non-interactive --username "__token__" --password "$PYPI_API_TOKEN" "$WHEELHOUSE/$wheel"
+    done
+}
+
+# Install twine if not already installed.
+python -m pip install --upgrade twine
+
+# Iterate over each project and upload its wheels.
+for project in "${PROJECTS[@]}"; do
+    upload_and_release_project $project
+done