added squeeze op

Author: mortont

lib/ortex/backend.ex

@@ -90,6 +90,23 @@ defmodule Ortex.Backend do
     put_in(out.data, %B{ref: Ortex.Native.reshape(ref, shape)})
   end
 
+  @impl true
+  def squeeze(out, tensor, axes) do
+    %T{shape: old_shape, names: names, data: %B{ref: ref}} = tensor
+    {new_shape, new_names} = Nx.Shape.squeeze(old_shape, axes, names)
+
+    if old_shape == new_shape do
+      %{out | data: %B{ref: ref}}
+    else
+      %{
+        out
+        | shape: new_shape,
+          names: new_names,
+          data: %B{ref: Ortex.Native.reshape(ref, new_shape |> Tuple.to_list())}
+      }
+    end
+  end
+
   if Application.compile_env(:ortex, :add_backend_on_inspect, true) do
     defp maybe_add_signature(result, %T{data: %B{ref: _mat_ref}}) do
       Inspect.Algebra.concat([

native/ortex/src/lib.rs

@@ -84,13 +84,15 @@ pub fn slice<'a>(
         strides,
     )))
 }
+
 #[rustler::nif]
 pub fn reshape<'a>(
     tensor: ResourceArc<OrtexTensor>,
     shape: Vec<usize>,
 ) -> NifResult<ResourceArc<OrtexTensor>> {
     Ok(ResourceArc::new(tensor.reshape(shape)?))
 }
+
 rustler::init!(
     "Elixir.Ortex.Native",
     [

test/shape/squeeze_test.exs

@@ -0,0 +1,39 @@
+defmodule Ortex.TestSqueeze do
+  use ExUnit.Case
+
+  test "1d squeeze" do
+    t = Nx.tensor([[[1, 2, 3, 4]]])
+    bin = t |> Nx.squeeze()
+
+    ort = t |> Nx.backend_copy(Ortex.Backend) |> Nx.squeeze() |> Nx.backend_transfer()
+
+    assert bin == ort
+  end
+
+  test "2d squeeze" do
+    t = Nx.tensor([[[[1, 2]], [[3, 4]]]])
+    bin = t |> Nx.squeeze()
+
+    ort = t |> Nx.backend_copy(Ortex.Backend) |> Nx.squeeze() |> Nx.backend_transfer()
+
+    assert bin == ort
+  end
+
+  test "axis squeeze" do
+    t = Nx.tensor([[[[1, 2]], [[3, 4]]]])
+    bin = t |> Nx.squeeze(axes: [0])
+
+    ort = t |> Nx.backend_copy(Ortex.Backend) |> Nx.squeeze(axes: [0]) |> Nx.backend_transfer()
+
+    assert bin == ort
+  end
+
+  test "named squeeze" do
+    t = Nx.tensor([[[[1, 2]], [[3, 4]]]], names: [:w, :x, :y, :z])
+    bin = t |> Nx.squeeze(axes: [:w])
+
+    ort = t |> Nx.backend_copy(Ortex.Backend) |> Nx.squeeze(axes: [:w]) |> Nx.backend_transfer()
+
+    assert bin == ort
+  end
+end