diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml
index b1a7b30..4af8bef 100644
--- a/.github/workflows/publish.yml
+++ b/.github/workflows/publish.yml
@@ -7,32 +7,68 @@ on:
   workflow_dispatch:
 
 jobs:
-  build-n-publish:
-    name: Publish to TestPyPI and PyPI
+  build:
+    name: Build distribution
     runs-on: ubuntu-latest
-    permissions:
-      id-token: write # For pypi trusted publishing
-
     steps:
       - uses: actions/checkout@v4
-
       - name: Set up Python
         uses: actions/setup-python@v5
         with:
           python-version: "3.x"
 
       - name: Install pypa/build
-        run: python -m pip install build --user
-
+        run: python3 -m pip install build --user
       - name: Build a binary wheel and a source tarball
-        run: python -m build --sdist --wheel --outdir dist/
+        run: python3 -m build
+      - name: Store the distribution packages
+        uses: actions/upload-artifact@v4
+        with:
+          name: python-package-distributions
+          path: dist/
+
+  publish-to-testpypi:
+    name: Publish to TestPyPI
+    needs:
+      - build
+    runs-on: ubuntu-latest
+
+    environment:
+      name: testpypi
+      url: https://test.pypi.org/p/flowjax
+
+    permissions:
+      id-token: write # IMPORTANT: mandatory for trusted publishing
 
-      - name: Publish distribution to Test PyPI
+    steps:
+      - name: Download all the dists
+        uses: actions/download-artifact@v4
+        with:
+          name: python-package-distributions
+          path: dist/
+      - name: Publish distribution 📦 to TestPyPI
         uses: pypa/gh-action-pypi-publish@release/v1
         with:
           repository-url: https://test.pypi.org/legacy/
 
-      - name: Publish distribution to PyPI
-        uses: pypa/gh-action-pypi-publish@release/v1
+  publish-to-pypi:
+    name: Publish Python distribution to PyPI
+    needs:
+      - publish-to-testpypi
+    runs-on: ubuntu-latest
+
+    environment:
+      name: pypi
+      url: https://pypi.org/p/flowjax
+
+    permissions:
+      id-token: write # IMPORTANT: mandatory for trusted publishing
+
+    steps:
+      - name: Download all the dists
+        uses: actions/download-artifact@v4
         with:
-          attestations: false # Newly introduced and seems to cause errors
+          name: python-package-distributions
+          path: dist/
+      - name: Publish distribution 📦 to PyPI
+        uses: pypa/gh-action-pypi-publish@release/v1