From ddbb293204e53eb2c3c91ad366b86506f7ca450f Mon Sep 17 00:00:00 2001
From: Joel Berkeley <16429957+joelberkeley@users.noreply.github.com>
Date: Sat, 30 Nov 2024 15:06:18 +0000
Subject: [PATCH 01/10] add support for Apple Metal

---
 .github/workflows/checks.yml                  | 59 +++++++++++++++++++
 .../apple-metal/PjrtPluginAppleMetal.idr      | 30 ++++++++++
 pjrt-plugins/apple-metal/README.md            | 10 ++++
 pjrt-plugins/apple-metal/build.sh             | 33 +++++++++++
 .../apple-metal/pjrt-plugin-apple-metal.ipkg  | 11 ++++
 pjrt-plugins/apple-metal/postinstall.sh       | 25 ++++++++
 .../XlaCuda.idr => apple-metal/Main.idr}      |  0
 test/apple-metal/apple-metal.ipkg             |  8 +++
 test/xla-cpu/{XlaCpu.idr => Main.idr}         |  0
 test/xla-cpu/xla-cpu.ipkg                     |  2 +-
 test/xla-cuda/Main.idr                        | 25 ++++++++
 test/xla-cuda/xla-cuda.ipkg                   |  2 +-
 12 files changed, 203 insertions(+), 2 deletions(-)
 create mode 100644 pjrt-plugins/apple-metal/PjrtPluginAppleMetal.idr
 create mode 100644 pjrt-plugins/apple-metal/README.md
 create mode 100644 pjrt-plugins/apple-metal/build.sh
 create mode 100644 pjrt-plugins/apple-metal/pjrt-plugin-apple-metal.ipkg
 create mode 100755 pjrt-plugins/apple-metal/postinstall.sh
 rename test/{xla-cuda/XlaCuda.idr => apple-metal/Main.idr} (100%)
 create mode 100644 test/apple-metal/apple-metal.ipkg
 rename test/xla-cpu/{XlaCpu.idr => Main.idr} (100%)
 create mode 100644 test/xla-cuda/Main.idr

diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml
index 30cc6c6be..6bb1ce585 100644
--- a/.github/workflows/checks.yml
+++ b/.github/workflows/checks.yml
@@ -112,6 +112,25 @@ jobs:
         name: pjrt_plugin_xla_cpu-darwin-aarch64
         path: pjrt_plugin_xla_cpu.dylib
         if-no-files-found: error
+  pjrt-plugin-apple-metal:
+    runs-on: macos-latest
+    steps:
+      - uses: actions/checkout@v4
+        with:
+          fetch-depth: 2
+      - name: Build or fetch Apple Metal PJRT plugin
+        run: |
+          prefix=jax_metal-0.1.0-py3-none-macosx_11_0_arm64
+          curl -fsL "https://files.pythonhosted.org/packages/80/af/ed482a421a868726e7ca3f51ac19b0c9a8e37f33f54413312c37e9056acc/jax_metal-0.1.0-py3-none-macosx_11_0_arm64.whl" \
+            -o "$prefix.zip"
+          unzip "$prefix.zip"
+          mv "$prefix/jax_plugins/metal_plugin/pjrt_plugin_metal_14.dylib" pjrt_plugin_apple_metal.dylib
+      - name: Upload binary
+        uses: actions/upload-artifact@v4
+        with:
+          name: pjrt_plugin_apple_metal
+          path: pjrt_plugin_apple_metal.dylib
+          if-no-files-found: error
   pjrt-plugin-xla-cuda-linux-x86_64:
     runs-on: ubuntu-latest
     steps:
@@ -181,6 +200,27 @@ jobs:
           name: tests-xla-cpu-darwin-aarch64
           path: test/xla-cpu/tests-xla-cpu.tar.gz
           if-no-files-found: error
+  build-tests-apple-metal:
+    runs-on: macos-latest
+    steps:
+    - uses: actions/checkout@v4
+    - name: Install build dependencies
+      run: |
+        brew install chezscheme
+        git clone https://github.com/stefan-hoeck/idris2-pack.git
+        (cd idris2-pack && make micropack SCHEME=chez)
+        ~/.pack/bin/pack switch HEAD
+    - name: Build tests
+      working-directory: test/xla-cpu
+      run: |
+        SPIDR_INSTALL_SUPPORT_LIBS=false ~/.pack/bin/pack --no-prompt build apple-metal.ipkg
+        tar cfz tests-apple-metal.tar.gz -C build/exec .
+    - name: Upload tests
+      uses: actions/upload-artifact@v4
+      with:
+        name: tests-apple-metal
+        path: test/xla-cpu/tests-apple-metal.tar.gz
+        if-no-files-found: error
   build-tests-xla-cuda-linux-x86_64:
     runs-on: ubuntu-latest
     container: ghcr.io/stefan-hoeck/idris2-pack
@@ -239,6 +279,25 @@ jobs:
       run: |
         tar xfz tests-xla-cpu.tar.gz && rm tests-xla-cpu.tar.gz
         ./test
+  test-apple-metal:
+    needs:
+      - pjrt-darwin-aarch64
+      - pjrt-plugin-apple-metal
+      - build-tests-apple-metal
+    runs-on: macos-latest
+    steps:
+    - name: Download artifacts
+      uses: actions/download-artifact@v4
+      with:
+        pattern: "{libc_xla-darwin-aarch64,pjrt_plugin_apple_metal,tests-apple-metal}"
+        merge-multiple: true
+    - name: Install runtime dependencies
+      run: |
+        brew install chezscheme
+    - name: Run tests
+      run: |
+        tar xfz tests-xla-cpu.tar.gz && rm tests-xla-cpu.tar.gz
+        ./test
   test-xla-cuda-linux-x86_64:
     needs:
       - pjrt-linux-x86_64
diff --git a/pjrt-plugins/apple-metal/PjrtPluginAppleMetal.idr b/pjrt-plugins/apple-metal/PjrtPluginAppleMetal.idr
new file mode 100644
index 000000000..a1182bdea
--- /dev/null
+++ b/pjrt-plugins/apple-metal/PjrtPluginAppleMetal.idr
@@ -0,0 +1,30 @@
+{--
+Copyright 2024 Joel Berkeley
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+--}
+module PjrtPluginAppleMEtal
+
+import System.FFI
+
+import public Compiler.Xla.PJRT.C.PjrtCApi
+import public Device
+
+%foreign "C:GetPjrtApi,pjrt_plugin_apple_metal"
+prim__getPjrtApi : PrimIO AnyPtr
+
+export
+device : Pjrt Device
+device = do
+  api <- MkPjrtApi <$> primIO prim__getPjrtApi
+  MkDevice api <$> pjrtClientCreate api
diff --git a/pjrt-plugins/apple-metal/README.md b/pjrt-plugins/apple-metal/README.md
new file mode 100644
index 000000000..8e3cc0c22
--- /dev/null
+++ b/pjrt-plugins/apple-metal/README.md
@@ -0,0 +1,10 @@
+# PJRT plugin for Apple Metal
+
+This is the PJRT plugin for Apple Metal, which provides hardware acceleration on Apple silicon.
+
+## Install
+
+On MacOS with AArch64 (ARM64, Apple Silicon), run
+```
+pack install pjrt-plugin-apple-metal
+```
diff --git a/pjrt-plugins/apple-metal/build.sh b/pjrt-plugins/apple-metal/build.sh
new file mode 100644
index 000000000..094d8c41d
--- /dev/null
+++ b/pjrt-plugins/apple-metal/build.sh
@@ -0,0 +1,33 @@
+#!/bin/sh -e
+
+script_dir=$(CDPATH="" cd -- "$(dirname -- "$0")" && pwd)
+cd "$script_dir/../.."
+. ./dev.sh
+rev=$(cat XLA_VERSION)
+
+osu="$(uname)"
+case $osu in
+  'Linux')
+    os=linux
+    arch=x86_64
+    ext=so
+    ;;
+  'Darwin')
+    os=darwin
+    arch=aarch64
+    ext=dylib
+    ;;
+  *)
+    echo "OS $osu not handled"
+    exit 1
+    ;;
+esac
+
+xla_dir=$(mktemp -d)
+install_xla "$rev" "$xla_dir"
+(
+  cd "$xla_dir"
+  ./configure.py --backend=CPU --os=$os
+  bazel build //xla/pjrt/c:pjrt_c_api_cpu_plugin.so
+)
+mv "$xla_dir/bazel-bin/xla/pjrt/c/pjrt_c_api_cpu_plugin.so" "pjrt_plugin_xla_cpu-$os-$arch.$ext"
diff --git a/pjrt-plugins/apple-metal/pjrt-plugin-apple-metal.ipkg b/pjrt-plugins/apple-metal/pjrt-plugin-apple-metal.ipkg
new file mode 100644
index 000000000..79183d236
--- /dev/null
+++ b/pjrt-plugins/apple-metal/pjrt-plugin-apple-metal.ipkg
@@ -0,0 +1,11 @@
+package pjrt-plugin-apple-metal
+version = 0.0.1
+
+depends = spidr
+modules = PjrtPluginAppleMetal
+
+brief = "XLA PJRT plugin for Apple Metal."
+readme = "README.md"
+license = "Apache License, Version 2.0"
+
+postinstall = "./postinstall.sh"
diff --git a/pjrt-plugins/apple-metal/postinstall.sh b/pjrt-plugins/apple-metal/postinstall.sh
new file mode 100755
index 000000000..7df8c4669
--- /dev/null
+++ b/pjrt-plugins/apple-metal/postinstall.sh
@@ -0,0 +1,25 @@
+#!/bin/sh -e
+
+if [ "$SPIDR_INSTALL_SUPPORT_LIBS" = false ]; then exit 0; fi
+
+script_dir=$(CDPATH="" cd -- "$(dirname -- "$0")" && pwd)
+cd "$script_dir/../.."
+
+os="$(uname)"
+case $os in
+  'Darwin')
+    ;;
+  *)
+    echo "WARNING: OS $os not supported, unable to fetch supporting libraries."
+    exit 0
+    ;;
+esac
+
+prefix=jax_metal-0.1.0-py3-none-macosx_11_0_arm64
+curl -fsL "https://files.pythonhosted.org/packages/80/af/ed482a421a868726e7ca3f51ac19b0c9a8e37f33f54413312c37e9056acc/jax_metal-0.1.0-py3-none-macosx_11_0_arm64.whl" \
+  -o "$prefix.zip"
+unzip "$prefix.zip"
+libdir="$(idris2 --libdir)/pjrt-plugin-xla-cpu-0.0.1/lib"
+mkdir -p libdir
+mv "$prefix/jax_plugins/metal_plugin/pjrt_plugin_metal_14.dylib" "$libdir/pjrt_plugin_apple_metal.dylib"
+rm -rf "$prefix.zip" $prefix
diff --git a/test/xla-cuda/XlaCuda.idr b/test/apple-metal/Main.idr
similarity index 100%
rename from test/xla-cuda/XlaCuda.idr
rename to test/apple-metal/Main.idr
diff --git a/test/apple-metal/apple-metal.ipkg b/test/apple-metal/apple-metal.ipkg
new file mode 100644
index 000000000..dd8efc992
--- /dev/null
+++ b/test/apple-metal/apple-metal.ipkg
@@ -0,0 +1,8 @@
+package apple-metal
+
+depends =
+  pjrt-plugin-apple-metal,
+  runner
+
+executable = test
+main = Main
diff --git a/test/xla-cpu/XlaCpu.idr b/test/xla-cpu/Main.idr
similarity index 100%
rename from test/xla-cpu/XlaCpu.idr
rename to test/xla-cpu/Main.idr
diff --git a/test/xla-cpu/xla-cpu.ipkg b/test/xla-cpu/xla-cpu.ipkg
index 24255b025..39fd35065 100644
--- a/test/xla-cpu/xla-cpu.ipkg
+++ b/test/xla-cpu/xla-cpu.ipkg
@@ -5,4 +5,4 @@ depends =
   runner
 
 executable = test
-main = XlaCpu
+main = Main
diff --git a/test/xla-cuda/Main.idr b/test/xla-cuda/Main.idr
new file mode 100644
index 000000000..422589049
--- /dev/null
+++ b/test/xla-cuda/Main.idr
@@ -0,0 +1,25 @@
+{--
+Copyright 2024 Joel Berkeley
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+--}
+module XlaCuda
+
+import System
+
+import TestRunner
+import PjrtPluginXlaCuda
+
+partial
+main : IO ()
+main = eitherT (die . show) run device
diff --git a/test/xla-cuda/xla-cuda.ipkg b/test/xla-cuda/xla-cuda.ipkg
index 66c3f269b..9d76e1994 100644
--- a/test/xla-cuda/xla-cuda.ipkg
+++ b/test/xla-cuda/xla-cuda.ipkg
@@ -5,4 +5,4 @@ depends =
   runner
 
 executable = test
-main = XlaCuda
+main = Main

From 48ddeb83ef765494705bbba4642172f5b6980523 Mon Sep 17 00:00:00 2001
From: Joel Berkeley <16429957+joelberkeley@users.noreply.github.com>
Date: Sat, 30 Nov 2024 15:07:53 +0000
Subject: [PATCH 02/10] ls

---
 .github/workflows/checks.yml | 1 +
 1 file changed, 1 insertion(+)

diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml
index 6bb1ce585..27fe6f9d1 100644
--- a/.github/workflows/checks.yml
+++ b/.github/workflows/checks.yml
@@ -124,6 +124,7 @@ jobs:
           curl -fsL "https://files.pythonhosted.org/packages/80/af/ed482a421a868726e7ca3f51ac19b0c9a8e37f33f54413312c37e9056acc/jax_metal-0.1.0-py3-none-macosx_11_0_arm64.whl" \
             -o "$prefix.zip"
           unzip "$prefix.zip"
+          ls
           mv "$prefix/jax_plugins/metal_plugin/pjrt_plugin_metal_14.dylib" pjrt_plugin_apple_metal.dylib
       - name: Upload binary
         uses: actions/upload-artifact@v4

From dd83193c53ac3a563c3b68343454e14a2368452c Mon Sep 17 00:00:00 2001
From: Joel Berkeley <16429957+joelberkeley@users.noreply.github.com>
Date: Sat, 30 Nov 2024 15:09:01 +0000
Subject: [PATCH 03/10] fix paths

---
 .github/workflows/checks.yml | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml
index 27fe6f9d1..a076195c3 100644
--- a/.github/workflows/checks.yml
+++ b/.github/workflows/checks.yml
@@ -124,8 +124,7 @@ jobs:
           curl -fsL "https://files.pythonhosted.org/packages/80/af/ed482a421a868726e7ca3f51ac19b0c9a8e37f33f54413312c37e9056acc/jax_metal-0.1.0-py3-none-macosx_11_0_arm64.whl" \
             -o "$prefix.zip"
           unzip "$prefix.zip"
-          ls
-          mv "$prefix/jax_plugins/metal_plugin/pjrt_plugin_metal_14.dylib" pjrt_plugin_apple_metal.dylib
+          mv "jax_plugins/metal_plugin/pjrt_plugin_metal_14.dylib" pjrt_plugin_apple_metal.dylib
       - name: Upload binary
         uses: actions/upload-artifact@v4
         with:

From 6730227dd3b4aa4b1a702a778ef89c79cd0ff539 Mon Sep 17 00:00:00 2001
From: Joel Berkeley <16429957+joelberkeley@users.noreply.github.com>
Date: Sat, 30 Nov 2024 15:14:47 +0000
Subject: [PATCH 04/10] names

---
 test/apple-metal/Main.idr | 4 ++--
 test/xla-cpu/Main.idr     | 2 +-
 test/xla-cuda/Main.idr    | 2 +-
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/test/apple-metal/Main.idr b/test/apple-metal/Main.idr
index 422589049..db388c772 100644
--- a/test/apple-metal/Main.idr
+++ b/test/apple-metal/Main.idr
@@ -13,12 +13,12 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 --}
-module XlaCuda
+module Main
 
 import System
 
 import TestRunner
-import PjrtPluginXlaCuda
+import PjrtPluginAppleMetal
 
 partial
 main : IO ()
diff --git a/test/xla-cpu/Main.idr b/test/xla-cpu/Main.idr
index 2e5d7b972..854d1eae4 100644
--- a/test/xla-cpu/Main.idr
+++ b/test/xla-cpu/Main.idr
@@ -13,7 +13,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 --}
-module XlaCpu
+module Main
 
 import System
 
diff --git a/test/xla-cuda/Main.idr b/test/xla-cuda/Main.idr
index 422589049..4a727f497 100644
--- a/test/xla-cuda/Main.idr
+++ b/test/xla-cuda/Main.idr
@@ -13,7 +13,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 --}
-module XlaCuda
+module Main
 
 import System
 

From d4943ddf57d11aa8b99204a69c73d5bd1a7346d7 Mon Sep 17 00:00:00 2001
From: Joel Berkeley <16429957+joelberkeley@users.noreply.github.com>
Date: Sat, 30 Nov 2024 16:01:41 +0000
Subject: [PATCH 05/10] dir

---
 .github/workflows/checks.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml
index a076195c3..070390b03 100644
--- a/.github/workflows/checks.yml
+++ b/.github/workflows/checks.yml
@@ -211,7 +211,7 @@ jobs:
         (cd idris2-pack && make micropack SCHEME=chez)
         ~/.pack/bin/pack switch HEAD
     - name: Build tests
-      working-directory: test/xla-cpu
+      working-directory: test/apple-metal
       run: |
         SPIDR_INSTALL_SUPPORT_LIBS=false ~/.pack/bin/pack --no-prompt build apple-metal.ipkg
         tar cfz tests-apple-metal.tar.gz -C build/exec .

From fa061e4db44118decf5054dd7a0b8c2289a9b5d2 Mon Sep 17 00:00:00 2001
From: Joel Berkeley <16429957+joelberkeley@users.noreply.github.com>
Date: Sat, 30 Nov 2024 16:18:35 +0000
Subject: [PATCH 06/10] pack.toml

---
 pack.toml | 5 +++++
 1 file changed, 5 insertions(+)

diff --git a/pack.toml b/pack.toml
index 7e1218ff7..ceb6eddd0 100644
--- a/pack.toml
+++ b/pack.toml
@@ -8,6 +8,11 @@ type = "local"
 path = ""
 ipkg = "test/runner/runner.ipkg"
 
+[custom.all.pjrt-plugin-apple-metal]
+type = "local"
+path = ""
+ipkg = "pjrt-plugins/xla-cpu/pjrt-plugin-apple-metal.ipkg"
+
 [custom.all.pjrt-plugin-xla-cpu]
 type = "local"
 path = ""

From 80dc785af0d0736c8fb883c03506c29db9d94428 Mon Sep 17 00:00:00 2001
From: Joel Berkeley <16429957+joelberkeley@users.noreply.github.com>
Date: Sat, 30 Nov 2024 16:35:02 +0000
Subject: [PATCH 07/10] wip

---
 .github/workflows/checks.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml
index 070390b03..b6331573e 100644
--- a/.github/workflows/checks.yml
+++ b/.github/workflows/checks.yml
@@ -219,7 +219,7 @@ jobs:
       uses: actions/upload-artifact@v4
       with:
         name: tests-apple-metal
-        path: test/xla-cpu/tests-apple-metal.tar.gz
+        path: test/apple-metal/tests-apple-metal.tar.gz
         if-no-files-found: error
   build-tests-xla-cuda-linux-x86_64:
     runs-on: ubuntu-latest

From bd23ffcf336096cffe55ead326776cb1f3640ccc Mon Sep 17 00:00:00 2001
From: Joel Berkeley <16429957+joelberkeley@users.noreply.github.com>
Date: Sat, 30 Nov 2024 16:54:29 +0000
Subject: [PATCH 08/10] wip

---
 pack.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pack.toml b/pack.toml
index ceb6eddd0..33edbc328 100644
--- a/pack.toml
+++ b/pack.toml
@@ -11,7 +11,7 @@ ipkg = "test/runner/runner.ipkg"
 [custom.all.pjrt-plugin-apple-metal]
 type = "local"
 path = ""
-ipkg = "pjrt-plugins/xla-cpu/pjrt-plugin-apple-metal.ipkg"
+ipkg = "pjrt-plugins/apple-metal/pjrt-plugin-apple-metal.ipkg"
 
 [custom.all.pjrt-plugin-xla-cpu]
 type = "local"

From 246e249367a4cc12738afc54dadcabd44d81250c Mon Sep 17 00:00:00 2001
From: Joel Berkeley <16429957+joelberkeley@users.noreply.github.com>
Date: Sat, 30 Nov 2024 17:29:05 +0000
Subject: [PATCH 09/10] e

---
 pjrt-plugins/apple-metal/PjrtPluginAppleMetal.idr | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pjrt-plugins/apple-metal/PjrtPluginAppleMetal.idr b/pjrt-plugins/apple-metal/PjrtPluginAppleMetal.idr
index a1182bdea..f9e1c33e6 100644
--- a/pjrt-plugins/apple-metal/PjrtPluginAppleMetal.idr
+++ b/pjrt-plugins/apple-metal/PjrtPluginAppleMetal.idr
@@ -13,7 +13,7 @@ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 See the License for the specific language governing permissions and
 limitations under the License.
 --}
-module PjrtPluginAppleMEtal
+module PjrtPluginAppleMetal
 
 import System.FFI
 

From f77bc606218cc8426e7abb72945f1f08d342e06b Mon Sep 17 00:00:00 2001
From: Joel Berkeley <16429957+joelberkeley@users.noreply.github.com>
Date: Sat, 30 Nov 2024 17:45:00 +0000
Subject: [PATCH 10/10] wip

---
 .github/workflows/checks.yml       | 2 +-
 pjrt-plugins/README.md             | 2 +-
 pjrt-plugins/apple-metal/README.md | 4 ++--
 3 files changed, 4 insertions(+), 4 deletions(-)

diff --git a/.github/workflows/checks.yml b/.github/workflows/checks.yml
index b6331573e..4a018d196 100644
--- a/.github/workflows/checks.yml
+++ b/.github/workflows/checks.yml
@@ -296,7 +296,7 @@ jobs:
         brew install chezscheme
     - name: Run tests
       run: |
-        tar xfz tests-xla-cpu.tar.gz && rm tests-xla-cpu.tar.gz
+        tar xfz tests-apple-metal.tar.gz && rm tests-apple-metal.tar.gz
         ./test
   test-xla-cuda-linux-x86_64:
     needs:
diff --git a/pjrt-plugins/README.md b/pjrt-plugins/README.md
index bc0d70f5f..8d977e52f 100644
--- a/pjrt-plugins/README.md
+++ b/pjrt-plugins/README.md
@@ -1,6 +1,6 @@
 # PJRT Plugins
 
-A PJRT plugin provides the compiler and hardware device support required to execute spidr graphs. We provide plugins for [CPU](xla-cpu/README.md) and [CUDA-enabled GPUs](xla-cuda/README.md). You can also use third-party plugins, or make your own.
+A PJRT plugin provides the compiler and hardware device support required to execute spidr graphs. We provide plugins for [CPU](xla-cpu/README.md), [Apple Metal](apple-metal/README.md), and [CUDA-enabled GPUs](xla-cuda/README.md). You can also use third-party plugins, or make your own.
 
 ## How to integrate your own plugin
 
diff --git a/pjrt-plugins/apple-metal/README.md b/pjrt-plugins/apple-metal/README.md
index 8e3cc0c22..fda57042c 100644
--- a/pjrt-plugins/apple-metal/README.md
+++ b/pjrt-plugins/apple-metal/README.md
@@ -1,10 +1,10 @@
 # PJRT plugin for Apple Metal
 
-This is the PJRT plugin for Apple Metal, which provides hardware acceleration on Apple silicon.
+This is the PJRT plugin for Apple Metal, which provides hardware acceleration with GPU on Apple silicon (AArch64, ARM64).
 
 ## Install
 
-On MacOS with AArch64 (ARM64, Apple Silicon), run
+Run
 ```
 pack install pjrt-plugin-apple-metal
 ```