fix: proper gpu support for linalg (#113)

* fix: add proper gpu support for linalg ops

* fix: no-op for same device

* chore: format

Author: polvalente

emlx/c_src/emlx_nif.cpp

@@ -495,6 +495,13 @@ NIF(eval) {
   return nx::nif::ok(env);
 }
 
+NIF(to_device) {
+  TENSOR_PARAM(0, t);
+  DEVICE_PARAM(1, device);
+  TENSOR(mlx::core::contiguous(*t, false, device));
+}
+ASYNC_NIF(to_device)
+
 NIF(stack) {
   LIST_PARAM(0, std::vector<mlx::core::array>, arrays);
   PARAM(1, int, axis);
@@ -1797,6 +1804,7 @@ ASYNC_NIF(window_scatter_min)
 
 static ErlNifFunc nif_funcs[] = {
     {"eval", 2, eval_async},
+    {"to_device", 3, to_device_async},
     {"to_blob", 2, to_blob_async},
     {"to_blob", 3, to_blob_async},
     {"tensor_to_shm", 3, tensor_to_shm_async},

emlx/lib/emlx.ex

@@ -1195,6 +1195,31 @@ defmodule EMLX do
     await_worker(job_ref)
   end
 
+  @doc """
+  Moves a tensor to `target_device` (`:cpu` or `:gpu`) without deallocating
+  the source. This is a no-op if the tensor is already on the target device.
+
+  Unlike `Nx.backend_transfer/2`, this does not round-trip through a binary
+  and does not call `backend_deallocate` on the original tensor. Internally
+  it schedules `mlx::core::contiguous(arr, target_device)` on the target
+  device's worker thread, which on Apple Silicon (unified memory) avoids any
+  physical data copy.
+  """
+  def to_device(tensor, target_device)
+
+  def to_device({dev, _} = tensor, dev), do: tensor
+
+  def to_device({old_device, ref} = tensor, target_device)
+      when is_tensor(old_device, ref) and target_device in [:cpu, :gpu] do
+    # Materialize the source on its own device's thread before constructing the
+    # contiguous op. Without this, evaluating contiguous(lazy_gpu_array, cpu)
+    # on the CPU worker would traverse the graph and hit a missing GPU stream.
+    eval(tensor)
+    {worker, effective_device} = resolve_worker(target_device)
+    job_ref = EMLX.NIF.to_device(worker, ref, effective_device) |> unwrap!()
+    await_worker(job_ref) |> wrap_tensor(effective_device)
+  end
+
   # ── Worker resolution ──────────────────────────────────────────────────────
   #
   # `Process.get(:emlx_command_queue)` is set by EMLX.CommandQueue.with_queue/2.

emlx/lib/emlx/backend.ex

@@ -175,14 +175,14 @@ defmodule EMLX.Backend do
     # Preserve quantization_config when copying a quantized tensor to EMLX.Backend.
     # The generic backend_copy goes through to_binary/from_binary which drops the config.
     target_device = device_option(opts)
-    gpu_opts = {EMLX.Backend, device: target_device}
+    device_opts = {EMLX.Backend, device: target_device}
 
     packed_size = Enum.reduce(Tuple.to_list(packed_shape), 1, &*/2)
     packed_binary = EMLX.to_blob(ref, packed_size)
     new_ref = EMLX.from_blob(packed_binary, packed_shape, :uint32, target_device)
 
-    new_scales = Nx.backend_transfer(cfg.scales, gpu_opts)
-    new_biases = Nx.backend_transfer(cfg.biases, gpu_opts)
+    new_scales = Nx.backend_copy(cfg.scales, device_opts)
+    new_biases = Nx.backend_copy(cfg.biases, device_opts)
 
     %T{
       shape: logical_shape,
@@ -1984,8 +1984,10 @@ defmodule EMLX.Backend do
       raise ArgumentError, "can't solve for singular matrix"
     end
 
-    a_typed = to_typed_ref(from_nx(a), a.type, {:f, 32})
-    b_mx = to_typed_ref(from_nx(b), b.type, {:f, 32})
+    {device, _} = from_nx(a)
+
+    a_typed = to_typed_ref(from_nx(a), a.type, {:f, 32}) |> EMLX.to_device(:cpu)
+    b_mx = to_typed_ref(from_nx(b), b.type, {:f, 32}) |> EMLX.to_device(:cpu)
 
     upper = !opts[:lower]
 
@@ -2018,6 +2020,7 @@ defmodule EMLX.Backend do
 
     out_mx
     |> EMLX.astype(to_mlx_type(out.type))
+    |> EMLX.to_device(device)
     |> to_nx(out)
   end
 
@@ -2103,11 +2106,18 @@ defmodule EMLX.Backend do
   end
 
   @impl true
-  def block(%Nx.Block.LinAlg.Cholesky{}, out, [tensor], _fun) do
-    from_nx(tensor)
-    |> to_typed_ref(tensor.type, {:f, 32})
-    |> EMLX.linalg_cholesky(false)
-    |> to_nx(out)
+  def block(%Nx.Block.LinAlg.Cholesky{} = struct, out, [tensor], fun) do
+    {device, _ref} = t_mx = from_nx(tensor)
+
+    case device do
+      :cpu ->
+        t = to_typed_ref(t_mx, tensor.type, {:f, 32})
+        cholesky = EMLX.linalg_cholesky(t, false)
+        to_nx(cholesky, out)
+
+      :gpu ->
+        fun.(struct, tensor)
+    end
   end
 
   @impl true
@@ -2120,10 +2130,18 @@ defmodule EMLX.Backend do
   end
 
   @impl true
-  def block(%Nx.Block.LinAlg.QR{mode: :reduced}, {out_q, out_r}, [tensor], _fun) do
-    t = to_typed_ref(from_nx(tensor), tensor.type, {:f, 32})
-    {q, r} = EMLX.linalg_qr(t)
-    {to_nx(q, out_q), to_nx(r, out_r)}
+  def block(%Nx.Block.LinAlg.QR{mode: :reduced} = struct, {out_q, out_r}, [tensor], fun) do
+    {device, _ref} = t_mx = from_nx(tensor)
+
+    case device do
+      :cpu ->
+        t = to_typed_ref(t_mx, tensor.type, {:f, 32})
+        {q, r} = EMLX.linalg_qr(t)
+        {to_nx(q, out_q), to_nx(r, out_r)}
+
+      :gpu ->
+        fun.(struct, tensor)
+    end
   end
 
   @impl true
@@ -2133,17 +2151,38 @@ defmodule EMLX.Backend do
   end
 
   @impl true
-  def block(%Nx.Block.LinAlg.Eigh{}, {out_eigenvals, out_eigenvecs}, [tensor], _fun) do
-    t = to_typed_ref(from_nx(tensor), tensor.type, {:f, 32})
-    {eigenvalues, eigenvectors} = EMLX.linalg_eigh(t, :L)
-    {to_nx(eigenvalues, out_eigenvals), to_nx(eigenvectors, out_eigenvecs)}
+  def block(%Nx.Block.LinAlg.Eigh{} = struct, {out_eigenvals, out_eigenvecs}, [tensor], fun) do
+    {device, _ref} = t_mx = from_nx(tensor)
+
+    case device do
+      :cpu ->
+        t = to_typed_ref(t_mx, tensor.type, {:f, 32})
+        {eigenvalues, eigenvectors} = EMLX.linalg_eigh(t, :L)
+        {to_nx(eigenvalues, out_eigenvals), to_nx(eigenvectors, out_eigenvecs)}
+
+      :gpu ->
+        fun.(struct, tensor)
+    end
   end
 
   @impl true
-  def block(%Nx.Block.LinAlg.SVD{full_matrices?: true}, {out_u, out_s, out_v}, [tensor], _fun) do
-    t = to_typed_ref(from_nx(tensor), tensor.type, {:f, 32})
-    [u, s, vt] = EMLX.linalg_svd(t, true)
-    {to_nx(u, out_u), to_nx(s, out_s), to_nx(vt, out_v)}
+  def block(
+        %Nx.Block.LinAlg.SVD{full_matrices?: true} = struct,
+        {out_u, out_s, out_v},
+        [tensor],
+        fun
+      ) do
+    {device, _ref} = t_mx = from_nx(tensor)
+
+    case device do
+      :cpu ->
+        t = to_typed_ref(t_mx, tensor.type, {:f, 32})
+        [u, s, vt] = EMLX.linalg_svd(t, true)
+        {to_nx(u, out_u), to_nx(s, out_s), to_nx(vt, out_v)}
+
+      :gpu ->
+        fun.(struct, tensor)
+    end
   end
 
   @impl true

emlx/lib/emlx/nif.ex

@@ -33,6 +33,10 @@ defmodule EMLX.NIF do
     :erlang.nif_error(:nif_not_loaded)
   end
 
+  def to_device(_worker, _tensor, _device) do
+    :erlang.nif_error(:nif_not_loaded)
+  end
+
   def item(_worker, _tensor_ref) do
     :erlang.nif_error(:nif_not_loaded)
   end

emlx/test/emlx/to_device_test.exs

@@ -0,0 +1,105 @@
+defmodule EMLX.ToDeviceTest do
+  use EMLX.Case, async: false
+
+  describe "CPU → CPU" do
+    test "preserves values" do
+      t = Nx.iota({3, 4}, type: :f32, backend: {EMLX.Backend, device: :cpu})
+      result_ref = EMLX.to_device(EMLX.Backend.from_nx(t), :cpu)
+
+      assert EMLX.to_blob(result_ref) ==
+               Nx.to_binary(Nx.iota({3, 4}, type: :f32, backend: Nx.BinaryBackend))
+    end
+
+    test "result is tagged :cpu" do
+      t = Nx.iota({4}, type: :f32, backend: {EMLX.Backend, device: :cpu})
+      result_ref = EMLX.to_device(EMLX.Backend.from_nx(t), :cpu)
+      assert elem(result_ref, 0) == :cpu
+    end
+
+    test "source tensor remains accessible after to_device" do
+      t = Nx.iota({4}, type: :f32, backend: {EMLX.Backend, device: :cpu})
+      ref = EMLX.Backend.from_nx(t)
+      _result = EMLX.to_device(ref, :cpu)
+      # Source must not be deallocated — to_blob would raise on deallocation
+      assert is_binary(EMLX.to_blob(ref))
+    end
+
+    test "works on a non-contiguous (transposed) tensor" do
+      t = Nx.iota({3, 4}, type: :f32, backend: {EMLX.Backend, device: :cpu}) |> Nx.transpose()
+      ref = EMLX.Backend.from_nx(t)
+      result_ref = EMLX.to_device(ref, :cpu)
+      assert elem(result_ref, 0) == :cpu
+      assert is_binary(EMLX.to_blob(result_ref))
+    end
+  end
+
+  describe "CPU → GPU" do
+    @tag :metal
+    test "result is tagged :gpu" do
+      t = Nx.iota({4}, type: :f32, backend: {EMLX.Backend, device: :cpu})
+      result_ref = EMLX.to_device(EMLX.Backend.from_nx(t), :gpu)
+      assert elem(result_ref, 0) == :gpu
+    end
+
+    @tag :metal
+    test "values are preserved" do
+      t = Nx.iota({3, 4}, type: :f32, backend: {EMLX.Backend, device: :cpu})
+      result_ref = EMLX.to_device(EMLX.Backend.from_nx(t), :gpu)
+      binary = EMLX.to_blob(result_ref)
+      expected = Nx.to_binary(Nx.iota({3, 4}, type: :f32, backend: Nx.BinaryBackend))
+      assert binary == expected
+    end
+
+    @tag :metal
+    test "source tensor remains accessible after transfer" do
+      t = Nx.iota({4}, type: :f32, backend: {EMLX.Backend, device: :cpu})
+      ref = EMLX.Backend.from_nx(t)
+      _result_ref = EMLX.to_device(ref, :gpu)
+      assert is_binary(EMLX.to_blob(ref))
+    end
+  end
+
+  describe "GPU → CPU" do
+    @tag :metal
+    test "result is tagged :cpu" do
+      t = Nx.iota({4}, type: :f32, backend: {EMLX.Backend, device: :gpu})
+      result_ref = EMLX.to_device(EMLX.Backend.from_nx(t), :cpu)
+      assert elem(result_ref, 0) == :cpu
+    end
+
+    @tag :metal
+    test "values are preserved" do
+      t = Nx.iota({3, 4}, type: :f32, backend: {EMLX.Backend, device: :gpu})
+      result_ref = EMLX.to_device(EMLX.Backend.from_nx(t), :cpu)
+      binary = EMLX.to_blob(result_ref)
+      expected = Nx.to_binary(Nx.iota({3, 4}, type: :f32, backend: Nx.BinaryBackend))
+      assert binary == expected
+    end
+
+    @tag :metal
+    test "source tensor remains accessible after transfer" do
+      t = Nx.iota({4}, type: :f32, backend: {EMLX.Backend, device: :gpu})
+      ref = EMLX.Backend.from_nx(t)
+      _result_ref = EMLX.to_device(ref, :cpu)
+      assert is_binary(EMLX.to_blob(ref))
+    end
+  end
+
+  describe "GPU → GPU" do
+    @tag :metal
+    test "result is tagged :gpu" do
+      t = Nx.iota({4}, type: :f32, backend: {EMLX.Backend, device: :gpu})
+      result_ref = EMLX.to_device(EMLX.Backend.from_nx(t), :gpu)
+      assert elem(result_ref, 0) == :gpu
+    end
+
+    @tag :metal
+    test "values are preserved" do
+      t = Nx.iota({3, 4}, type: :f32, backend: {EMLX.Backend, device: :gpu})
+      result_ref = EMLX.to_device(EMLX.Backend.from_nx(t), :gpu)
+      binary = EMLX.to_blob(result_ref)
+      expected = Nx.to_binary(Nx.iota({3, 4}, type: :f32, backend: Nx.BinaryBackend))
+      assert binary == expected
+    end
+  end
+end

emlx_axon/mix.exs

@@ -1,7 +1,7 @@
 defmodule EMLXAxon.MixProject do
   use Mix.Project
 
-  @version "0.1.0"
+  @version "0.3.0"
   @source_url "https://github.com/elixir-nx/emlx"
 
   def project do
@@ -24,11 +24,10 @@ defmodule EMLXAxon.MixProject do
 
   defp deps do
     [
-      {:emlx, path: "../emlx"},
+      # {:emlx, path: "../emlx"},
+      {:emlx, "~> 0.3"},
       {:axon, "~> 0.7"},
-      # Inherit emlx's Nx git pin to avoid constraint conflicts.
-      {:nx, github: "elixir-nx/nx", sparse: "nx", override: true},
-      {:bumblebee, github: "elixir-nx/bumblebee", override: true, only: [:dev, :test]},
+      {:bumblebee, "~> 0.7"},
       {:ex_doc, "~> 0.34", only: :docs}
     ]
   end

emlx_axon/mix.lock

@@ -1,18 +1,19 @@
 %{
-  "axon": {:hex, :axon, "0.7.0", "2e2c6d93b4afcfa812566b8922204fa022b60081e86ebd411df4db7ea30f5457", [:mix], [{:kino, "~> 0.7", [hex: :kino, repo: "hexpm", optional: true]}, {:kino_vega_lite, "~> 0.1.7", [hex: :kino_vega_lite, repo: "hexpm", optional: true]}, {:nx, "~> 0.9", [hex: :nx, repo: "hexpm", optional: false]}, {:polaris, "~> 0.1", [hex: :polaris, repo: "hexpm", optional: false]}, {:table_rex, "~> 3.1.1", [hex: :table_rex, repo: "hexpm", optional: true]}], "hexpm", "ee9857a143c9486597ceff434e6ca833dc1241be6158b01025b8217757ed1036"},
-  "bumblebee": {:git, "https://github.com/elixir-nx/bumblebee.git", "0fd8114cf5429af9236f100f3350986e9d823c02", []},
+  "axon": {:hex, :axon, "0.8.1", "c4a975e62a14ab6c374997b77367ec3a4c2740952ac474d3b0f202c91b7f75c4", [:mix], [{:kino, "~> 0.7", [hex: :kino, repo: "hexpm", optional: true]}, {:kino_vega_lite, "~> 0.1.7", [hex: :kino_vega_lite, repo: "hexpm", optional: true]}, {:nx, "~> 0.10", [hex: :nx, repo: "hexpm", optional: false]}, {:polaris, "~> 0.1", [hex: :polaris, repo: "hexpm", optional: false]}, {:table_rex, "~> 3.1 or ~> 4.1", [hex: :table_rex, repo: "hexpm", optional: true]}], "hexpm", "682a3517489300507ac9345f28341e7fa95bc5b4960d645816074ce551795d37"},
+  "bumblebee": {:hex, :bumblebee, "0.7.0", "3532b1487e6c4f5f75189f7b820508f51fb7159238c68f97286427f87dc99654", [:mix], [{:axon, "~> 0.8.0", [hex: :axon, repo: "hexpm", optional: false]}, {:jason, "~> 1.4.0", [hex: :jason, repo: "hexpm", optional: false]}, {:nx, "~> 0.12.0", [hex: :nx, repo: "hexpm", optional: false]}, {:nx_image, "~> 0.1.0", [hex: :nx_image, repo: "hexpm", optional: false]}, {:nx_signal, "~> 0.2.0", [hex: :nx_signal, repo: "hexpm", optional: false]}, {:progress_bar, "~> 3.0", [hex: :progress_bar, repo: "hexpm", optional: false]}, {:safetensors, "~> 0.1.3", [hex: :safetensors, repo: "hexpm", optional: false]}, {:tokenizers, "~> 0.4", [hex: :tokenizers, repo: "hexpm", optional: false]}, {:unpickler, "~> 0.1.0", [hex: :unpickler, repo: "hexpm", optional: false]}, {:unzip, "~> 0.12.0 or ~> 0.13.0", [hex: :unzip, repo: "hexpm", optional: false]}], "hexpm", "40244e9e98227777084c44e7a2e3a7908bcddb02d5bb852bbb615163d5d076c9"},
   "castore": {:hex, :castore, "1.0.18", "5e43ef0ec7d31195dfa5a65a86e6131db999d074179d2ba5a8de11fe14570f55", [:mix], [], "hexpm", "f393e4fe6317829b158fb74d86eb681f737d2fe326aa61ccf6293c4104957e34"},
-  "complex": {:hex, :complex, "0.6.0", "b0130086a7a8c33574d293b2e0e250f4685580418eac52a5658a4bd148f3ccf1", [:mix], [], "hexpm", "0a5fa95580dcaf30fcd60fe1aaf24327c0fe401e98c24d892e172e79498269f9"},
+  "complex": {:hex, :complex, "0.7.0", "695632ef9487517aa5d57edd1697801079d622414cb2e1a7cf538b1f9a50f205", [:mix], [], "hexpm", "0ee39c0803129f546e7f3f640da8f021c9e659402bf59da6f7f2c4848f068f8d"},
   "decimal": {:hex, :decimal, "2.4.1", "6c0fbede12fb122ba685e9ab41c6a40c129e322b3aa192f9e072e61f3a6ffaf2", [:mix], [], "hexpm", "7e618897933a8455f19a727d7c5e50a2c071a544b700e5e724298ecb4340187f"},
   "earmark_parser": {:hex, :earmark_parser, "1.4.44", "f20830dd6b5c77afe2b063777ddbbff09f9759396500cdbe7523efd58d7a339c", [:mix], [], "hexpm", "4778ac752b4701a5599215f7030989c989ffdc4f6df457c5f36938cc2d2a2750"},
   "elixir_make": {:hex, :elixir_make, "0.9.0", "6484b3cd8c0cee58f09f05ecaf1a140a8c97670671a6a0e7ab4dc326c3109726", [:mix], [], "hexpm", "db23d4fd8b757462ad02f8aa73431a426fe6671c80b200d9710caf3d1dd0ffdb"},
+  "emlx": {:hex, :emlx, "0.3.0", "71fa22717d4a768ed2353f052845b2cbd6840c9e50dc67b1c9ccd5731d54ecf8", [:make, :mix], [{:elixir_make, "~> 0.6", [hex: :elixir_make, repo: "hexpm", optional: false]}, {:nx, "~> 0.12", [hex: :nx, repo: "hexpm", optional: false]}], "hexpm", "2778038adb05c4a2813a24eb5a89112c41a0b5b00350f059b348af0e01891c58"},
   "ex_doc": {:hex, :ex_doc, "0.40.2", "f50edec428c4b0a457a167de42414c461122a3585a99515a69d09fff19e5597e", [:mix], [{:earmark_parser, "~> 1.4.44", [hex: :earmark_parser, repo: "hexpm", optional: false]}, {:makeup_c, ">= 0.1.0", [hex: :makeup_c, repo: "hexpm", optional: true]}, {:makeup_elixir, "~> 0.14 or ~> 1.0", [hex: :makeup_elixir, repo: "hexpm", optional: false]}, {:makeup_erlang, "~> 0.1 or ~> 1.0", [hex: :makeup_erlang, repo: "hexpm", optional: false]}, {:makeup_html, ">= 0.1.0", [hex: :makeup_html, repo: "hexpm", optional: true]}], "hexpm", "4fa426e2beb47854a162e2c488727fdec51cd4692e319b23810c2804cb1a40fe"},
   "jason": {:hex, :jason, "1.4.5", "2e3a008590b0b8d7388c20293e9dcc9cf3e5d642fd2a114e4cbbb52e595d940a", [:mix], [{:decimal, "~> 1.0 or ~> 2.0 or ~> 3.0", [hex: :decimal, repo: "hexpm", optional: true]}], "hexpm", "b0c823996102bcd0239b3c2444eb00409b72f6a140c1950bc8b457d836b30684"},
   "makeup": {:hex, :makeup, "1.2.1", "e90ac1c65589ef354378def3ba19d401e739ee7ee06fb47f94c687016e3713d1", [:mix], [{:nimble_parsec, "~> 1.4", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "d36484867b0bae0fea568d10131197a4c2e47056a6fbe84922bf6ba71c8d17ce"},
   "makeup_elixir": {:hex, :makeup_elixir, "1.0.1", "e928a4f984e795e41e3abd27bfc09f51db16ab8ba1aebdba2b3a575437efafc2", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}, {:nimble_parsec, "~> 1.2.3 or ~> 1.3", [hex: :nimble_parsec, repo: "hexpm", optional: false]}], "hexpm", "7284900d412a3e5cfd97fdaed4f5ed389b8f2b4cb49efc0eb3bd10e2febf9507"},
   "makeup_erlang": {:hex, :makeup_erlang, "1.1.0", "835f7e60792e08824cda445639555d7bf1bbbddb1b60b306e33cb6f6db24dc74", [:mix], [{:makeup, "~> 1.0", [hex: :makeup, repo: "hexpm", optional: false]}], "hexpm", "1cd6780fb1dd1a03979abaed0fe82712b0625118fd5257d3ebbf73f960c73c3c"},
   "nimble_parsec": {:hex, :nimble_parsec, "1.4.2", "8efba0122db06df95bfaa78f791344a89352ba04baedd3849593bfce4d0dc1c6", [:mix], [], "hexpm", "4b21398942dda052b403bbe1da991ccd03a053668d147d53fb8c4e0efe09c973"},
-  "nx": {:git, "https://github.com/elixir-nx/nx.git", "1c3d86eb635e136eab7d5a1a12a57794c200d204", [sparse: "nx"]},
+  "nx": {:hex, :nx, "0.12.0", "32bc205bab5486d73892132d17a11ea113e97427a29bb70606a544724b95e193", [:mix], [{:complex, "~> 0.7", [hex: :complex, repo: "hexpm", optional: false]}, {:telemetry, "~> 0.4.0 or ~> 1.0", [hex: :telemetry, repo: "hexpm", optional: false]}], "hexpm", "7d022a33ea3c900eb6e2e91b4e0793759459c886f482be61978004b5e4843b5e"},
   "nx_image": {:hex, :nx_image, "0.1.2", "0c6e3453c1dc30fc80c723a54861204304cebc8a89ed3b806b972c73ee5d119d", [:mix], [{:nx, "~> 0.4", [hex: :nx, repo: "hexpm", optional: false]}], "hexpm", "9161863c42405ddccb6dbbbeae078ad23e30201509cc804b3b3a7c9e98764b81"},
   "nx_signal": {:hex, :nx_signal, "0.2.0", "e1ca0318877b17c81ce8906329f5125f1e2361e4c4235a5baac8a95ee88ea98e", [:mix], [{:nx, "~> 0.6", [hex: :nx, repo: "hexpm", optional: false]}], "hexpm", "7247e5e18a177a59c4cb5355952900c62fdeadeb2bad02a9a34237b68744e2bb"},
   "polaris": {:hex, :polaris, "0.1.0", "dca61b18e3e801ecdae6ac9f0eca5f19792b44a5cb4b8d63db50fc40fc038d22", [:mix], [{:nx, "~> 0.5", [hex: :nx, repo: "hexpm", optional: false]}], "hexpm", "13ef2b166650e533cb24b10e2f3b8ab4f2f449ba4d63156e8c569527f206e2c2"},

emlx_axon/test/test_helper.exs

@@ -1 +1,24 @@
-ExUnit.start(exclude: [:bumblebee, :metal])
+use_gpu? =
+  String.downcase(System.get_env("EMLX_TEST_DEFAULT_GPU", "false")) in [
+    "1",
+    "true",
+    "yes",
+    "t",
+    "y"
+  ]
+
+if use_gpu? do
+  Application.put_env(:nx, :default_backend, {EMLX.Backend, device: :gpu})
+end
+
+gpu_exclude =
+  if use_gpu? do
+    case EMLX.NIF.command_queue_new(:gpu) do
+      {:ok, _} -> []
+      {:error, _} -> [:metal]
+    end
+  else
+    [:metal]
+  end
+
+ExUnit.start(exclude: [:bumblebee] ++ gpu_exclude)