Document take as optional

elixir-nx · May 12, 2024 · dde9152 · dde9152
1 parent b8f2254
commit dde9152
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/nx/lib/nx/backend.ex b/nx/lib/nx/backend.ex
@@ -158,6 +158,7 @@ defmodule Nx.Backend do
 
   @callback all_close(out :: tensor, tensor, tensor, keyword) :: tensor
   @callback top_k(out :: tensor, tensor, keyword) :: tensor
+  @callback take(out :: tensor, input :: tensor, indices :: tensor, keyword) :: tensor
 
   @optional_callbacks [
     optional: 3,
@@ -176,7 +177,8 @@ defmodule Nx.Backend do
     ifft2: 3,
     qr: 3,
     cholesky: 2,
-    eigh: 3
+    eigh: 3,
+    take: 4
   ]
 
   ## Inspect implementation