Does EXLA support the Apple Silicon GPU (MPS)?

Hi!,

I have the Mac Studio and I want to use the GPU on it with EXLA.
I made the project that has config like below.

import Config
config :nx, :default_backend, EXLA.Backend

with iex ,
Nx.tensor [1,2,3,4] ,
it show
[info] TfrtCpuClient created.
#Nx.Tensor<
s64[4]
EXLA.Backend<host:0, 0.1189406505.872022046.161358>
[1, 2, 3, 4]

it seems not use GPU.

And, with below code (to make long Nx calculation) ----
url =
“https://pix4free.org/assets/library/2021-01-12/originals/san_francisco_california_golden_gate_bridge_water.jpg”

%{body: raw_image} = Req.get!(url)

imageNx = StbImage.read_binary!(raw_image)
            |> StbImage.to_nx()
            |> Nx.reshape({:auto,3})

IO.puts "Started --------"
start_ts =DateTime.utc_now() |> DateTime.to_unix(:second)

KMeans.fit(imageNx,
    num_clusters: 10,
    num_runs: 10,
    max_iterations: 200,
    key: Nx.Random.key(42)
  )

end_ts =DateTime.utc_now() |> DateTime.to_unix(:second)
IO.puts "DONE -- ellapse : #{end_ts - start_ts}"

It takes about 230 seconds, and the Mac Activity Monitor show 0 GPU usage of beam_smp.

Does the EXLA support Apple Silicon GPU ?

Any suggestion will be appreciated..

Thanks,