actual reshape implementation

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,12 +81,13 @@ 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, _tensor) do
-    out
+  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
 
   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,24 @@ 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

@@ -1,4 +1,5 @@
 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