Merge pull request #12 from elixir-nx/tm-backend-ops

`reshape` and `squeeze` implementations on `Ortex.Backend`

Author: mortont

lib/ortex/backend.ex

@@ -71,7 +71,6 @@ defmodule Ortex.Backend do
   @impl true
   def inspect(%T{} = tensor, inspect_opts) do
     limit = if inspect_opts.limit == :infinity, do: :infinity, else: inspect_opts.limit + 1
-    # IO.inspect(limit)
 
     tensor
     |> to_binary(min(limit, Nx.size(tensor)))
@@ -82,7 +81,30 @@ defmodule Ortex.Backend do
   @impl true
   def slice(out, %T{data: %B{ref: tensor_ref}}, start_indicies, lengths, strides) do
     r = Ortex.Native.slice(tensor_ref, start_indicies, lengths, strides)
-    put_in(out.data, %Ortex.Backend{ref: r})
+    put_in(out.data, %B{ref: r})
+  end
+
+  @impl true
+  def reshape(out, %T{data: %B{ref: ref}}) do
+    shape = Nx.shape(out) |> Tuple.to_list()
+    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

lib/ortex/native.ex

@@ -24,4 +24,6 @@ defmodule Ortex.Native do
 
   def slice(_tensor, _start_indicies, _lengths, _strides),
     do: :erlang.nif_error(:nif_not_loaded)
+
+  def reshape(_tensor, _shape), do: :erlang.nif_error(:nif_not_loaded)
 end

native/ortex/src/lib.rs

@@ -84,9 +84,26 @@ 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",
-    [run, init, from_binary, to_binary, show_session, slice],
+    [
+        run,
+        init,
+        from_binary,
+        to_binary,
+        show_session,
+        slice,
+        reshape
+    ],
     load = |env: Env, _| {
         rustler::resource!(OrtexModel, env);
         rustler::resource!(OrtexTensor, env);

native/ortex/src/tensor.rs

@@ -63,6 +63,71 @@ impl OrtexTensor {
         }
     }
 
+    pub fn reshape(&self, shape: Vec<usize>) -> rustler::NifResult<Self> {
+        match self {
+            OrtexTensor::s8(y) => {
+                Ok(OrtexTensor::s8(y.clone().into_shape(shape).map_err(
+                    |e| rustler::Error::Term(Box::new(e.to_string())),
+                )?))
+            }
+            OrtexTensor::s16(y) => {
+                Ok(OrtexTensor::s16(y.clone().into_shape(shape).map_err(
+                    |e| rustler::Error::Term(Box::new(e.to_string())),
+                )?))
+            }
+            OrtexTensor::s32(y) => {
+                Ok(OrtexTensor::s32(y.clone().into_shape(shape).map_err(
+                    |e| rustler::Error::Term(Box::new(e.to_string())),
+                )?))
+            }
+            OrtexTensor::s64(y) => {
+                Ok(OrtexTensor::s64(y.clone().into_shape(shape).map_err(
+                    |e| rustler::Error::Term(Box::new(e.to_string())),
+                )?))
+            }
+            OrtexTensor::u8(y) => {
+                Ok(OrtexTensor::u8(y.clone().into_shape(shape).map_err(
+                    |e| rustler::Error::Term(Box::new(e.to_string())),
+                )?))
+            }
+            OrtexTensor::u16(y) => {
+                Ok(OrtexTensor::u16(y.clone().into_shape(shape).map_err(
+                    |e| rustler::Error::Term(Box::new(e.to_string())),
+                )?))
+            }
+            OrtexTensor::u32(y) => {
+                Ok(OrtexTensor::u32(y.clone().into_shape(shape).map_err(
+                    |e| rustler::Error::Term(Box::new(e.to_string())),
+                )?))
+            }
+            OrtexTensor::u64(y) => {
+                Ok(OrtexTensor::u64(y.clone().into_shape(shape).map_err(
+                    |e| rustler::Error::Term(Box::new(e.to_string())),
+                )?))
+            }
+            OrtexTensor::f16(y) => {
+                Ok(OrtexTensor::f16(y.clone().into_shape(shape).map_err(
+                    |e| rustler::Error::Term(Box::new(e.to_string())),
+                )?))
+            }
+            OrtexTensor::bf16(y) => {
+                Ok(OrtexTensor::bf16(y.clone().into_shape(shape).map_err(
+                    |e| rustler::Error::Term(Box::new(e.to_string())),
+                )?))
+            }
+            OrtexTensor::f32(y) => {
+                Ok(OrtexTensor::f32(y.clone().into_shape(shape).map_err(
+                    |e| rustler::Error::Term(Box::new(e.to_string())),
+                )?))
+            }
+            OrtexTensor::f64(y) => {
+                Ok(OrtexTensor::f64(y.clone().into_shape(shape).map_err(
+                    |e| rustler::Error::Term(Box::new(e.to_string())),
+                )?))
+            }
+        }
+    }
+
     pub fn dtype(&self) -> (Atom, usize) {
         match self {
             OrtexTensor::s8(_) => (ortex_atoms::s(), 8),

test/shape/reshape_test.exs

@@ -0,0 +1,55 @@
+defmodule Ortex.TestReshape do
+  # TODO: Fix this, it is not truly validating the reshaping on the ortex side
+  use ExUnit.Case
+
+  test "1d reshape" do
+    t = Nx.tensor([1, 2, 3, 4])
+    bin = t |> Nx.reshape({2, 2})
+
+    ort = t |> Nx.backend_copy(Ortex.Backend) |> Nx.reshape({2, 2}) |> Nx.backend_transfer()
+
+    assert bin == ort
+  end
+
+  test "2d reshape" do
+    shape = Nx.tensor([[0], [0], [0], [0]])
+    t = Nx.tensor([1, 2, 3, 4])
+    bin = t |> Nx.reshape(shape)
+
+    ort =
+      t
+      |> Nx.backend_copy(Ortex.Backend)
+      |> Nx.reshape(shape |> Nx.backend_copy(Ortex.Backend))
+      |> Nx.backend_transfer()
+
+    assert bin == ort
+  end
+
+  test "scalar reshape" do
+    shape = {1, 1, 1}
+    t = Nx.tensor(1)
+    bin = t |> Nx.reshape(shape)
+
+    ort =
+      t
+      |> Nx.backend_copy(Ortex.Backend)
+      |> Nx.reshape(shape)
+      |> Nx.backend_transfer()
+
+    assert bin == ort
+  end
+
+  test "auto reshape" do
+    shape = {:auto, 2}
+    t = Nx.tensor([[1, 2, 3], [4, 5, 6]])
+    bin = t |> Nx.reshape(shape)
+
+    ort =
+      t
+      |> Nx.backend_copy(Ortex.Backend)
+      |> Nx.reshape(shape)
+      |> Nx.backend_transfer()
+
+    assert bin == ort
+  end
+end

test/slice/slice_test.exs test/shape/slice_test.exstest/slice/slice_test.exs renamed to test/shape/slice_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