add make file

.github/workflows/lint.yml

@@ -12,6 +12,11 @@ on:
       - "docs/**"
       - "**.md"
 
+# Cancel in-progress runs when a new commit is pushed
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
+
 jobs:
   lint:
     name: Lint with Ruff
@@ -22,22 +27,11 @@ jobs:
       - name: Install uv
         uses: astral-sh/setup-uv@v5
         with:
-          version: "0.9.5"
+          enable-cache: true
+          cache-dependency-glob: "pyproject.toml"
 
       - name: Set up Python
-        uses: actions/setup-python@v6
-        with:
-          python-version: "3.11"
-
-      - name: Install dependencies
-        run: |
-          uv pip install --system ruff
-
-      - name: Run ruff check
-        run: |
-          ruff check torchax test test_dist
-
-      - name: Run ruff format check
-        run: |
-          ruff format --check torchax test test_dist
+        run: uv python install 3.11
 
+      - name: Run linter
+        run: make lint-check

.github/workflows/nightly.yml

@@ -19,27 +19,32 @@ jobs:
       - name: Checkout code
         uses: actions/checkout@v4
 
-      - name: Set up Python
-        uses: actions/setup-python@v5
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
         with:
-          python-version: '3.11'
+          enable-cache: true
+
+      - name: Set up Python
+        run: uv python install 3.11
 
-      - name: Install build dependencies
-        run: pip install build
+      - name: Create virtual environment
+        run: uv venv
 
       - name: Patch version for Nightly
         run: |
+          source .venv/bin/activate
           python scripts/update_to_nightly_version.py
-
           # Verify the change (Debugging)
           grep "__version__" torchax/__init__.py
 
-      - name: Build Wheel and Sdist
-        run: python -m build --wheel
+      - name: Build package
+        run: |
+          source .venv/bin/activate
+          make build
 
       - name: Publish to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1
-        # Note: We are NOT providing a password here because we will use Trusted Publishing.
+        # Note: Using Trusted Publishing (passwordless)
         # If you must use a token, uncomment the lines below:
         # with:
         #   password: ${{ secrets.PYPI_API_TOKEN }}

.github/workflows/release.yml

@@ -16,7 +16,7 @@ jobs:
     # in PyPI Settings > Publishing for your release configuration.
     environment:
       name: pypi
-      url: https://pypi.org/p/torchax # Replace 'torchax' with your package name
+      url: https://pypi.org/p/torchax
 
     permissions:
       # MANDATORY for OIDC Trusted Publishing
@@ -30,18 +30,22 @@ jobs:
         with:
           fetch-depth: 0
 
-      - name: Set up Python
-        uses: actions/setup-python@v5
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
         with:
-          python-version: '3.11'
+          enable-cache: true
+
+      - name: Set up Python
+        run: uv python install 3.11
 
-      - name: Install build dependencies
-        run: pip install build
+      - name: Create virtual environment
+        run: uv venv
 
-      - name: Build Wheel
-        # We build both here, as stable releases usually offer both formats.
-        run: python -m build --wheel
+      - name: Build package
+        run: |
+          source .venv/bin/activate
+          make build
 
       - name: Publish to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1
-        # Trusted Publishing handles authentication automatically.
+        # Trusted Publishing handles authentication automatically

.github/workflows/static.yml

@@ -31,43 +31,49 @@ jobs:
     steps:
       - name: Checkout
         uses: actions/checkout@v5
-      - name: Set up Python 3.12
-        uses: actions/setup-python@v6
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
         with:
-          python-version: 3.12
+          enable-cache: true
+          cache-dependency-glob: "pyproject.toml"
+
+      - name: Set up Python
+        run: uv python install 3.12
+
+      - name: Create virtual environment
+        run: uv venv
+
       - name: Setup Pages
         uses: actions/configure-pages@v5
+
       - name: Install dependencies
         run: |
-          pip install --upgrade pip
-          # note: installing torch and torchvision together ensures their version compatibility
-          pip install torch==2.8.0 torchvision --index-url https://download.pytorch.org/whl/cpu
-          pip install jupyter nbconvert pandas matplotlib
-          pip install -r test-requirements.txt
-          pip install -e .[cpu]
-          pip install mkdocs mkdocs-material mkdocs-rtd-dropdown mkdocs-jupyter
-      - name: build docs
-        working-directory: docs
-        run: |
-          mkdocs build
+          source .venv/bin/activate
+          uv pip install -e ".[cpu,docs]"
+
       - name: Generate notebooks
         working-directory: docs
         run: |
+          source ../.venv/bin/activate
           export JAX_PLATFORMS=cpu
           for FILE in $(find docs -name '*.py'); do
             export OUTFILE="${FILE%.*}.ipynb"
-            jupytext --to ipynb $FILE -o "$OUTFILE"
-            jupyter nbconvert --to notebook --execute $OUTFILE --output $(basename $OUTFILE)
+            uv tool run jupytext --to ipynb $FILE -o "$OUTFILE"
+            uv tool run jupyter nbconvert --to notebook --execute $OUTFILE --output $(basename $OUTFILE)
           done
-      - name: Build site
-        working-directory: docs
+
+      - name: Build documentation
         run: |
-          mkdocs build
+          source .venv/bin/activate
+          make docs
+
       - name: Upload artifact
         uses: actions/upload-pages-artifact@v3
         with:
-          # Upload entire repository
+          # Upload docs site
           path: "./docs/site"
+
       - name: Deploy to GitHub Pages
         id: deployment
         uses: actions/deploy-pages@v4

.github/workflows/unit-tests.yml

@@ -1,21 +1,28 @@
-# This workflow will install Python dependencies, run tests and lint with a variety of Python versions
+# This workflow will install Python dependencies and run tests with multiple Python versions
 # For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python
 
-name: Python package
+name: Python unit tests
 
 on:
   push:
     branches: ["main"]
     paths-ignore:
       - "docs/**"
+      - "**.md"
   pull_request:
     branches: ["main"]
     paths-ignore:
       - "docs/**"
+      - "**.md"
+
+# Cancel in-progress runs when a new commit is pushed
+concurrency:
+  group: ${{ github.workflow }}-${{ github.ref }}
+  cancel-in-progress: true
 
 jobs:
-  build:
-    name: Python unit tests
+  test:
+    name: Python unit tests (${{ matrix.python-version }})
     runs-on: ubuntu-latest
     strategy:
       fail-fast: false
@@ -25,39 +32,41 @@ jobs:
     steps:
       - uses: actions/checkout@v5
 
-      - name: Set up Python ${{ matrix.python-version }}
-        uses: actions/setup-python@v6
+      - name: Install uv
+        uses: astral-sh/setup-uv@v5
         with:
-          python-version: ${{ matrix.python-version }}
+          enable-cache: true
+          cache-dependency-glob: "pyproject.toml"
+
+      - name: Set up Python ${{ matrix.python-version }}
+        run: uv python install ${{ matrix.python-version }}
+
+      - name: Create virtual environment
+        run: uv venv --python ${{ matrix.python-version }}
+
       - name: Install dependencies
         run: |
-          pip install --upgrade pip
-          # note: installing torch and torchvision together ensures their version compatibility
-          pip install torch==2.8.0 torchvision --index-url https://download.pytorch.org/whl/cpu
-          pip install jupyter nbconvert pandas matplotlib
-          pip install -r test-requirements.txt
-          pip install -e .[cpu]
-      - name: Test with pytest
-        shell: bash
+          source .venv/bin/activate
+          make install-test
+
+      - name: Run unit tests
+        run: |
+          source .venv/bin/activate
+          make test
+        continue-on-error: false
+
+      - name: Run distributed tests
         run: |
-          export JAX_PLATFORMS=cpu
-          # Find all Python test files recursively
-          find ./test -name "test_*.py" -type f | while IFS= read -r test_file; do
-            # Skip tests with known issues
-            if [[ "$test_file" == *"test_tf_integration.py"* ]]; then
-              echo "Skipping ${test_file}. TODO(https://github.com/pytorch/xla/issues/8770): Investigate"
-              continue
-            fi
-            echo "Running tests for $test_file"
-            pytest "$test_file"
-          done
-          # Run distributed tests.
-          XLA_FLAGS=--xla_force_host_platform_device_count=4 pytest -n 0 test_dist/
-          echo "Tests completed."
+          source .venv/bin/activate
+          XLA_FLAGS=--xla_force_host_platform_device_count=4 pytest test_dist/ -n 0
+
       - name: Test tutorials can run
-        shell: bash
         run: |
+          source .venv/bin/activate
           export JAX_PLATFORMS=cpu
-          for FILE in $(find docs -name '*.py'); do
-              python $FILE
+          for FILE in $(find docs/docs/tutorials -name '*.py' 2>/dev/null || true); do
+            if [ -f "$FILE" ]; then
+              echo "Testing $FILE"
+              python "$FILE"
+            fi
           done

.gitignore

@@ -153,6 +153,7 @@ activemq-data/
 .env
 .envrc
 .venv
+.local
 env/
 venv/
 ENV/

.python-version

@@ -0,0 +1,2 @@
+3.11
+