Only slice again if there are strides

elixir-nx · Nov 5, 2023 · 3783da6 · 3783da6
1 parent 1fe44ce
commit 3783da6
Showing 1 changed file with 12 additions and 2 deletions.
diff --git a/exla/lib/exla/defn.ex b/exla/lib/exla/defn.ex
@@ -1401,7 +1401,12 @@ defmodule EXLA.Defn do
     else
       zeros = List.duplicate(0, tuple_size(ans.shape))
       slice = Value.dynamic_slice(tensor, start_indices, lengths)
-      Value.slice(slice, zeros, lengths, strides)
+
+      if Enum.all?(strides, & &1 == 1) do
+        slice
+      else
+        Value.slice(slice, zeros, lengths, strides)
+      end
     end
   end
 
@@ -1414,7 +1419,12 @@ defmodule EXLA.Defn do
     else
       zeros = List.duplicate(0, tuple_size(ans.shape))
       slice = EXLA.Op.dynamic_slice(tensor, start_indices, lengths)
-      EXLA.Op.slice(slice, zeros, lengths, strides)
+
+      if Enum.all?(strides, & &1 == 1) do
+        slice
+      else
+        EXLA.Op.slice(slice, zeros, lengths, strides)
+      end
     end
   end