diff --git a/.github/workflows/delete-docker-on-pr-close.yml b/.github/workflows/delete-docker-on-pr-close.yml
new file mode 100644
index 0000000..2a6d526
--- /dev/null
+++ b/.github/workflows/delete-docker-on-pr-close.yml
@@ -0,0 +1,155 @@
+name: Docker
+
+# This workflow uses actions that are not certified by GitHub.
+# They are provided by a third-party and are governed by
+# separate terms of service, privacy policy, and support
+# documentation.
+
+on:
+  pull_request:
+    branches: [ "main" ]
+    types: [ closed ]
+
+env:
+  # Use docker.io for Docker Hub if empty
+  REGISTRY: ghcr.io
+  JAX_VERSION : 0.4.23
+  JAX_CONDA_IMAGE_NAME : "yfukai/conda-jax"
+  # github.repository as <account>/<repo>
+  IMAGE_NAME: ${{ github.repository }}
+
+jobs:
+  build:
+
+    runs-on: ubuntu-latest
+    permissions:
+      contents: read
+      packages: write
+      # This is used to complete the identity challenge
+      # with sigstore/fulcio when running outside of PRs.
+      id-token: write
+
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v3
+
+      # Install the cosign tool except on PR
+      # https://github.com/sigstore/cosign-installer
+      - name: Install Cosign
+#        if: github.event_name != 'pull_request'
+        uses: sigstore/cosign-installer@v3.1.1
+        with:
+          cosign-release: 'v2.2.1'
+
+      # Workaround: https://github.com/docker/build-push-action/issues/461
+      - name: Setup Docker buildx
+        uses: docker/setup-buildx-action@79abd3f86f79a9d68a23c75a09a9a85889262adf
+
+      # Login against a Docker registry except on PR
+      # https://github.com/docker/login-action
+      - name: Log into registry ${{ env.REGISTRY }}
+#        if: github.event_name != 'pull_request'
+        uses: docker/login-action@28218f9b04b4f3f62068d7b6ce6ca5b26e35336c
+        with:
+          registry: ${{ env.REGISTRY }}
+          username: ${{ github.actor }}
+          password: ${{ secrets.GITHUB_TOKEN }}
+
+      # Extract metadata (tags, labels) for Docker
+      # https://github.com/docker/metadata-action
+      - name: Extract Docker metadata (JAX conda CPU)
+        id: meta_jax_conda
+        uses: docker/metadata-action@98669ae865ea3cffbcbaa878cf57c20bbf1c6c38
+        with:
+          images: ${{ env.REGISTRY }}/${{ env.JAX_CONDA_IMAGE_NAME }}
+          tags: |
+            type=schedule
+            type=ref,event=branch
+            type=ref,event=tag
+            type=ref,event=pr
+
+      - name: Extract Docker metadata (JAX conda GPU)
+        id: meta_jax_conda_cuda
+        uses: docker/metadata-action@98669ae865ea3cffbcbaa878cf57c20bbf1c6c38
+        with:
+          images: ${{ env.REGISTRY }}/${{ env.JAX_CONDA_IMAGE_NAME }}
+          tags: |
+            type=schedule,suffix=-cuda
+            type=ref,event=branch,suffix=-cuda
+            type=ref,event=tag,suffix=-cuda
+            type=ref,event=pr,suffix=-cuda
+
+      - name: Extract Docker metadata (CPU)
+        id: meta
+        uses: docker/metadata-action@98669ae865ea3cffbcbaa878cf57c20bbf1c6c38
+        with:
+          images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
+          tags: |
+            type=schedule
+            type=ref,event=branch
+            type=ref,event=tag
+            type=ref,event=pr
+
+      - name: Extract Docker metadata (GPU)
+        id: meta_cuda
+        uses: docker/metadata-action@98669ae865ea3cffbcbaa878cf57c20bbf1c6c38
+        with:
+          images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
+          tags: |
+            type=schedule,suffix=-cuda
+            type=ref,event=branch,suffix=-cuda
+            type=ref,event=tag,suffix=-cuda
+            type=ref,event=pr,suffix=-cuda
+
+#      # Build and push Docker image with Buildx
+#      # https://github.com/docker/build-push-action
+#      - name: Build and push Docker image (JAX conda CPU)
+#        id: build-and-push-jax-conda-cpu
+#        uses: docker/build-push-action@v5
+#        with:
+#          context: jax_image/
+#          platforms: linux/amd64,linux/arm64/v8
+#          push: true #${{ github.event_name != 'pull_request' }}
+#          tags: ${{ steps.meta_jax_conda.outputs.tags }}
+#          labels: ${{ steps.meta_jax_conda.outputs.labels }}
+#          build-args: |
+#            JAX_VERSION=${{ env.JAX_VERSION }}
+#
+#      - name: Build and push Docker image (JAX conda GPU)
+#        id: build-and-push-jax-conda-gpu
+#        uses: docker/build-push-action@v5
+#        with:
+#          context: jax_image/
+#          platforms: linux/amd64
+#          push: true #${{ github.event_name != 'pull_request' }}
+#          tags: ${{ steps.meta_jax_conda_cuda.outputs.tags }}
+#          labels: ${{ steps.meta_jax_conda_cuda.outputs.labels }}
+#          build-args: |
+#            JAX_VERSION=${{ env.JAX_VERSION }}
+#            JAX_VERSION_EXTRA=cuda11_local
+#            BASE_IMAGE=nvidia/cuda:11.8.0-cudnn8-devel-ubuntu20.04
+#
+#      - name: Build and push Docker image (CPU)
+#        id: build-and-push-cpu
+#        uses: docker/build-push-action@v5
+#        with:
+#          context: .
+#          platforms: linux/amd64,linux/arm64/v8
+#          push: true
+#          tags: ${{ steps.meta.outputs.tags }}
+#          labels: ${{ steps.meta.outputs.labels }}
+#          build-args: |
+#            BASE_IMAGE=${{ env.REGISTRY }}/${{ env.JAX_CONDA_IMAGE_NAME }}:${{ env.meta_jax_conda_first_tag }}
+#
+#      - name: Build and push Docker image (GPU)
+#        id: build-and-push-gpu
+#        uses: docker/build-push-action@v5
+#        with:
+#          context: .
+#          platforms: linux/amd64
+#          push: true
+#          tags: ${{ steps.meta_cuda.outputs.tags }}
+#          labels: ${{ steps.meta_cuda.outputs.labels }}
+#          build-args: |
+#            BASE_IMAGE=${{ env.REGISTRY }}/${{ env.JAX_CONDA_IMAGE_NAME }}:${{ env.meta_jax_conda_cuda_first_tag }}
+#
diff --git a/.github/workflows/docker-publish.yml b/.github/workflows/docker-publish.yml
index fbb635a..727cd53 100644
--- a/.github/workflows/docker-publish.yml
+++ b/.github/workflows/docker-publish.yml
@@ -15,8 +15,8 @@ env:
   # Use docker.io for Docker Hub if empty
   REGISTRY: ghcr.io
   JAX_VERSION : 0.4.23
-  # github.repository as <account>/<repo>
   JAX_CONDA_IMAGE_NAME : "yfukai/conda-jax"
+  # github.repository as <account>/<repo>
   IMAGE_NAME: ${{ github.repository }}
 
 jobs:
@@ -58,61 +58,108 @@ jobs:
 
       # Extract metadata (tags, labels) for Docker
       # https://github.com/docker/metadata-action
-      - name: Extract Docker metadata (JAX conda)
+      - name: Extract Docker metadata (JAX conda CPU)
         id: meta_jax_conda
         uses: docker/metadata-action@98669ae865ea3cffbcbaa878cf57c20bbf1c6c38
         with:
           images: ${{ env.REGISTRY }}/${{ env.JAX_CONDA_IMAGE_NAME }}
-
-      - name: Extract Docker metadata
+          tags: |
+            type=schedule
+            type=ref,event=branch
+            type=ref,event=tag
+            type=ref,event=pr
+
+      - name: Extract first tag (JAX conda CPU)
+        id: extract_first_tag_cpu
+        run: |
+          # Extracting the first component
+          first_component=$(echo "${{ steps.meta_jax_conda.outputs.tags }}" | cut -d ',' -f 1)
+          # Setting the output
+          echo "meta_jax_conda_first_tag=$first_component" >> $GITHUB_OUTPUT
+
+      - name: Extract Docker metadata (JAX conda GPU)
+        id: meta_jax_conda_cuda
+        uses: docker/metadata-action@98669ae865ea3cffbcbaa878cf57c20bbf1c6c38
+        with:
+          images: ${{ env.REGISTRY }}/${{ env.JAX_CONDA_IMAGE_NAME }}
+          tags: |
+            type=schedule,suffix=-cuda
+            type=ref,event=branch,suffix=-cuda
+            type=ref,event=tag,suffix=-cuda
+            type=ref,event=pr,suffix=-cuda
+
+      - name: Extract first tag (JAX conda GPU)
+        id: extract_first_tag_gpu
+        run: |
+          # Extracting the first component
+          first_component=$(echo "${{ steps.meta_jax_conda_cuda.outputs.tags }}" | cut -d ',' -f 1)
+          # Setting the output
+          echo "meta_jax_conda_cuda_first_tag=$first_component" >> $GITHUB_OUTPUT
+
+      - name: Extract Docker metadata (CPU)
         id: meta
         uses: docker/metadata-action@98669ae865ea3cffbcbaa878cf57c20bbf1c6c38
         with:
           images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
-
+          tags: |
+            type=schedule
+            type=ref,event=branch
+            type=ref,event=tag
+            type=ref,event=pr
+
+      - name: Extract Docker metadata (GPU)
+        id: meta_cuda
+        uses: docker/metadata-action@98669ae865ea3cffbcbaa878cf57c20bbf1c6c38
+        with:
+          images: ${{ env.REGISTRY }}/${{ env.IMAGE_NAME }}
+          tags: |
+            type=schedule,suffix=-cuda
+            type=ref,event=branch,suffix=-cuda
+            type=ref,event=tag,suffix=-cuda
+            type=ref,event=pr,suffix=-cuda
 
       - name: Set up QEMU
         uses: docker/setup-qemu-action@v2
 
-      # Build and push Docker image with Buildx (don't push on PR)
+      # Build and push Docker image with Buildx
       # https://github.com/docker/build-push-action
-#      - name: Build and push Docker image (JAX conda CPU)
-#        id: build-and-push-jax-conda-cpu
-#        uses: docker/build-push-action@v5
-#        with:
-#          context: jax_image/
-#          platforms: linux/amd64,linux/arm64/v8
-#          push: true #${{ github.event_name != 'pull_request' }}
-#          tags: ${{ steps.meta_jax_conda.outputs.tags }}
-#          labels: ${{ steps.meta_jax_conda.outputs.labels }}
-#          build-args: |
-#            JAX_VERSION=${{ env.JAX_VERSION }}
-
-#      - name: Build and push Docker image (JAX conda GPU)
-#        id: build-and-push-jax-conda-gpu
-#        uses: docker/build-push-action@v5
-#        with:
-#          context: jax_image/
-#          platforms: linux/amd64
-#          push: true #${{ github.event_name != 'pull_request' }}
-#          tags: ${{ steps.meta_jax_conda.outputs.tags }}-cuda
-#          labels: ${{ steps.meta_jax_conda.outputs.labels }}
-#          build-args: |
-#            JAX_VERSION=${{ env.JAX_VERSION }}
-#            JAX_VERSION_EXTRA=cuda11_local
-#            BASE_IMAGE=nvidia/cuda:11.8.0-cudnn8-devel-ubuntu20.04
-#
+      - name: Build and push Docker image (JAX conda CPU)
+        id: build-and-push-jax-conda-cpu
+        uses: docker/build-push-action@v5
+        with:
+          context: jax_image/
+          platforms: linux/amd64,linux/arm64/v8
+          push: true
+          tags: ${{ steps.meta_jax_conda.outputs.tags }}
+          labels: ${{ steps.meta_jax_conda.outputs.labels }}
+          build-args: |
+            JAX_VERSION=${{ env.JAX_VERSION }}
+
+      - name: Build and push Docker image (JAX conda GPU)
+        id: build-and-push-jax-conda-gpu
+        uses: docker/build-push-action@v5
+        with:
+          context: jax_image/
+          platforms: linux/amd64
+          push: true
+          tags: ${{ steps.meta_jax_conda_cuda.outputs.tags }}
+          labels: ${{ steps.meta_jax_conda_cuda.outputs.labels }}
+          build-args: |
+            JAX_VERSION=${{ env.JAX_VERSION }}
+            JAX_VERSION_EXTRA=cuda11_local
+            BASE_IMAGE=nvidia/cuda:11.8.0-cudnn8-devel-ubuntu20.04
+
       - name: Build and push Docker image (CPU)
         id: build-and-push-cpu
         uses: docker/build-push-action@v5
         with:
           context: .
           platforms: linux/amd64,linux/arm64/v8
-          push: true # ${{ github.event_name != 'pull_request' }}
+          push: true
           tags: ${{ steps.meta.outputs.tags }}
           labels: ${{ steps.meta.outputs.labels }}
           build-args: |
-            BASE_IMAGE=${{ env.REGISTRY }}/${{ env.JAX_CONDA_IMAGE_NAME }}:pr-8
+            BASE_IMAGE=${{ env.REGISTRY }}/${{ env.JAX_CONDA_IMAGE_NAME }}:${{ env.meta_jax_conda_first_tag }}
 
       - name: Build and push Docker image (GPU)
         id: build-and-push-gpu
@@ -120,8 +167,8 @@ jobs:
         with:
           context: .
           platforms: linux/amd64
-          push: true #${{ github.event_name != 'pull_request' }}
-          tags: ${{ steps.meta.outputs.tags }}-cuda
-          labels: ${{ steps.meta.outputs.labels }}
+          push: true
+          tags: ${{ steps.meta_cuda.outputs.tags }}
+          labels: ${{ steps.meta_cuda.outputs.labels }}
           build-args: |
-            BASE_IMAGE=${{ env.REGISTRY }}/${{ env.JAX_CONDA_IMAGE_NAME }}:pr-8-cuda
+            BASE_IMAGE=${{ env.REGISTRY }}/${{ env.JAX_CONDA_IMAGE_NAME }}:${{ env.meta_jax_conda_cuda_first_tag }}