Merge pull request #9 from relaypro-open/develop

implemented `slice` on `Ortex.Backend`

Author: mortont

lib/ortex/backend.ex

@@ -79,6 +79,12 @@ defmodule Ortex.Backend do
     |> maybe_add_signature(tensor)
   end
 
+  @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})
+  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([

lib/ortex/native.ex

@@ -21,4 +21,7 @@ defmodule Ortex.Native do
   def from_binary(_bin, _shape, _type), do: :erlang.nif_error(:nif_not_loaded)
   def to_binary(_reference, _bits, _limit), do: :erlang.nif_error(:nif_not_loaded)
   def show_session(_model), do: :erlang.nif_error(:nif_not_loaded)
+
+  def slice(_tensor, _start_indicies, _lengths, _strides),
+    do: :erlang.nif_error(:nif_not_loaded)
 end

native/ortex/src/lib.rs

@@ -71,9 +71,22 @@ fn to_binary<'a>(
     utils::to_binary(env, reference, bits, limit)
 }
 
+#[rustler::nif]
+pub fn slice<'a>(
+    tensor: ResourceArc<OrtexTensor>,
+    start_indicies: Vec<isize>,
+    lengths: Vec<isize>,
+    strides: Vec<isize>,
+) -> NifResult<ResourceArc<OrtexTensor>> {
+    Ok(ResourceArc::new(tensor.slice(
+        start_indicies,
+        lengths,
+        strides,
+    )))
+}
 rustler::init!(
     "Elixir.Ortex.Native",
-    [run, init, from_binary, to_binary, show_session],
+    [run, init, from_binary, to_binary, show_session, slice],
     load = |env: Env, _| {
         rustler::resource!(OrtexModel, env);
         rustler::resource!(OrtexTensor, env);

native/ortex/src/tensor.rs

@@ -1,6 +1,6 @@
 //! Conversions for packing/unpacking `OrtexTensor`s into different types
 use ndarray::prelude::*;
-use ndarray::Data;
+use ndarray::{ArrayBase, ArrayView, Data, IxDyn};
 use ort::tensor::{DynOrtTensor, FromArray, InputTensor, TensorElementDataType};
 use ort::OrtError;
 use rustler::Atom;
@@ -9,7 +9,8 @@ use crate::constants::ortex_atoms;
 
 #[derive(Debug)]
 #[allow(non_camel_case_types)]
-/// Enum for wrapping different types to pass back to the BEAM
+/// Enum for wrapping different types to pass back to the BEAM since rustler can't
+/// pass type generics back and forth
 pub enum OrtexTensor {
     s8(Array<i8, IxDyn>),
     s16(Array<i16, IxDyn>),
@@ -96,6 +97,49 @@ impl OrtexTensor {
         };
         contents
     }
+
+    pub fn slice<'a>(
+        &'a self,
+        start_indicies: Vec<isize>,
+        lengths: Vec<isize>,
+        strides: Vec<isize>,
+    ) -> Self {
+        let mut slice_specs: Vec<(isize, Option<isize>, isize)> = vec![];
+        for ((start_index, length), stride) in start_indicies
+            .iter()
+            .zip(lengths.iter())
+            .zip(strides.iter())
+        {
+            slice_specs.push((*start_index, Some(*length + *start_index), *stride));
+        }
+        match self {
+            OrtexTensor::s8(y) => OrtexTensor::s8(slice_array(y, &slice_specs).to_owned()),
+            OrtexTensor::s16(y) => OrtexTensor::s16(slice_array(y, &slice_specs).to_owned()),
+            OrtexTensor::s32(y) => OrtexTensor::s32(slice_array(y, &slice_specs).to_owned()),
+            OrtexTensor::s64(y) => OrtexTensor::s64(slice_array(y, &slice_specs).to_owned()),
+            OrtexTensor::u8(y) => OrtexTensor::u8(slice_array(y, &slice_specs).to_owned()),
+            OrtexTensor::u16(y) => OrtexTensor::u16(slice_array(y, &slice_specs).to_owned()),
+            OrtexTensor::u32(y) => OrtexTensor::u32(slice_array(y, &slice_specs).to_owned()),
+            OrtexTensor::u64(y) => OrtexTensor::u64(slice_array(y, &slice_specs).to_owned()),
+            OrtexTensor::f16(y) => OrtexTensor::f16(slice_array(y, &slice_specs).to_owned()),
+            OrtexTensor::bf16(y) => OrtexTensor::bf16(slice_array(y, &slice_specs).to_owned()),
+            OrtexTensor::f32(y) => OrtexTensor::f32(slice_array(y, &slice_specs).to_owned()),
+            OrtexTensor::f64(y) => OrtexTensor::f64(slice_array(y, &slice_specs).to_owned()),
+        }
+    }
+}
+
+fn slice_array<'a, T, D>(
+    array: &'a Array<T, D>,
+    slice_specs: &'a Vec<(isize, Option<isize>, isize)>,
+) -> ArrayView<'a, T, D>
+where
+    D: Dimension,
+{
+    array.slice_each_axis(|ax: ndarray::AxisDescription| {
+        let (start, end, step) = slice_specs[ax.axis.index()];
+        ndarray::Slice { start, end, step }
+    })
 }
 
 fn get_bytes<'a, T>(array: &'a ArrayBase<T, IxDyn>) -> &'a [u8]

test/slice/slice_test.exs

@@ -0,0 +1,97 @@
+defmodule Ortex.TestSlice do
+  use ExUnit.Case
+
+  {tensor1d, _} = Nx.Random.uniform(Nx.Random.key(42), 0, 256, shape: {10})
+  {tensor2d, _} = Nx.Random.uniform(Nx.Random.key(42), 0, 256, shape: {10, 10})
+
+  @tensor1d tensor1d
+  @tensor2d tensor2d
+
+  defp tensor_binary(tensor, dtype) do
+    tensor |> Nx.as_type(dtype)
+  end
+
+  defp tensor_ortex(tensor, dtype) do
+    tensor
+    |> Nx.as_type(dtype)
+    |> Nx.backend_transfer(Ortex.Backend)
+  end
+
+  test "1d slice f32" do
+    bin = tensor_binary(@tensor1d, :f32) |> Nx.slice([0], [4])
+
+    ort = tensor_ortex(@tensor1d, :f32) |> Nx.slice([0], [4]) |> Nx.backend_transfer()
+
+    assert bin == ort
+  end
+
+  test "2d slice f32" do
+    bin = tensor_binary(@tensor2d, :f32) |> Nx.slice([0, 2], [4, 6])
+
+    ort =
+      tensor_ortex(@tensor2d, :f32)
+      |> Nx.slice([0, 2], [4, 6])
+      |> Nx.backend_transfer()
+
+    assert bin == ort
+  end
+
+  test "1d slice u8" do
+    bin = tensor_binary(@tensor1d, :u8) |> Nx.slice([0], [4])
+
+    ort = tensor_ortex(@tensor1d, :u8) |> Nx.slice([0], [4]) |> Nx.backend_transfer()
+
+    assert bin == ort
+  end
+
+  test "2d slice u8" do
+    bin = tensor_binary(@tensor2d, :u8) |> Nx.slice([0, 2], [4, 6])
+
+    ort =
+      tensor_ortex(@tensor2d, :u8)
+      |> Nx.slice([0, 2], [4, 6])
+      |> Nx.backend_transfer()
+
+    assert bin == ort
+  end
+
+  test "1d slice f32 strided" do
+    bin = tensor_binary(@tensor1d, :f32) |> Nx.slice([0], [4], strides: [2])
+
+    ort =
+      tensor_ortex(@tensor1d, :f32) |> Nx.slice([0], [4], strides: [2]) |> Nx.backend_transfer()
+
+    assert bin == ort
+  end
+
+  test "2d slice f32 strided" do
+    bin = tensor_binary(@tensor2d, :f32) |> Nx.slice([0, 2], [4, 6], strides: [2, 1])
+
+    ort =
+      tensor_ortex(@tensor2d, :f32)
+      |> Nx.slice([0, 2], [4, 6], strides: [2, 1])
+      |> Nx.backend_transfer()
+
+    assert bin == ort
+  end
+
+  test "1d slice u8 strided" do
+    bin = tensor_binary(@tensor1d, :u8) |> Nx.slice([0], [4], strides: [2])
+
+    ort =
+      tensor_ortex(@tensor1d, :u8) |> Nx.slice([0], [4], strides: [2]) |> Nx.backend_transfer()
+
+    assert bin == ort
+  end
+
+  test "2d slice u8 strided" do
+    bin = tensor_binary(@tensor2d, :u8) |> Nx.slice([0, 2], [4, 6], strides: [2, 1])
+
+    ort =
+      tensor_ortex(@tensor2d, :u8)
+      |> Nx.slice([0, 2], [4, 6], strides: [2, 1])
+      |> Nx.backend_transfer()
+
+    assert bin == ort
+  end
+end