Merge pull request #87 from andfoy/add_vstack

Add vstack

lib/extorch/native/tensor/ops/manipulation.ex

@@ -946,5 +946,54 @@ defmodule ExTorch.Native.Tensor.Ops.Manipulation do
     """
     @spec permute(ExTorch.Tensor.t(), tuple() | [integer()]) :: ExTorch.Tensor.t()
     defbinding(permute(input, dims))
+
+    @doc """
+    Stack tensors in sequence vertically (row wise).
+
+    This is equivalent to concatenation along the first axis after all 1-D tensors
+    have been reshaped by `ExTorch.atleast_2d/1`.
+
+    ## Arguments
+    - `tensors` (`[ExTorch.Tensor.t()] | tuple()`) - sequence of tensors to concatenate.
+
+    ## Optional arguments
+     - out (`ExTorch.Tensor | nil`) - an optional pre-allocated tensor used to
+    store the output result. Default: `nil`
+
+    ## Examples
+        iex> a = ExTorch.tensor([1, 2, 3])
+        iex> b = ExTorch.tensor([4, 5, 6])
+        iex> ExTorch.vstack({a, b})
+        #Tensor<
+        [[1, 2, 3],
+         [4, 5, 6]]
+        [size: {2, 3}, dtype: :byte, device: :cpu, requires_grad: false]>
+
+        iex> a = ExTorch.tensor([[1],[2],[3]])
+        #Tensor<
+        [[1],
+         [2],
+         [3]]
+        [size: {3, 1}, dtype: :byte, device: :cpu, requires_grad: false]>
+        iex> b = ExTorch.tensor([[4],[5],[6]])
+        #Tensor<
+        [[4],
+         [5],
+         [6]]
+        [size: {3, 1}, dtype: :byte, device: :cpu, requires_grad: false]>
+
+        iex> ExTorch.vstack([a, b])
+        #Tensor<
+        [[1],
+         [2],
+         [3],
+         [4],
+         [5],
+         [6]]
+        [size: {6, 1}, dtype: :byte, device: :cpu, requires_grad: false]>
+
+    """
+    @spec vstack([ExTorch.Tensor.t()] | tuple(), ExTorch.Tensor.t() | nil) :: ExTorch.Tensor.t()
+    defbinding(vstack(tensors, out \\ nil), fn_aliases: [:row_stack])
   end
 end

native/extorch/include/manipulation.h

@@ -109,3 +109,5 @@ TensorTuple nonzero(
 std::shared_ptr<CrossTensor> permute(
     const std::shared_ptr<CrossTensor> &input,
     rust::Vec<int64_t> dims);
+
+std::shared_ptr<CrossTensor> vstack(TensorList tensor_list, TensorOut opt_out);

native/extorch/src/csrc/manipulation.cc

@@ -501,3 +501,17 @@ std::shared_ptr<CrossTensor> permute(
 
     return std::make_shared<CrossTensor>(std::move(out_tensor));
 }
+
+std::shared_ptr<CrossTensor> vstack(TensorList tensor_list, TensorOut opt_out) {
+    std::vector<CrossTensor> tensor_vec = unpack_tensor_list(tensor_list);
+    CrossTensor out_tensor;
+
+    if(opt_out.used) {
+        out_tensor = *opt_out.tensor.get();
+        out_tensor = torch::vstack_out(out_tensor, tensor_vec);
+    } else {
+        out_tensor = torch::vstack(tensor_vec);
+    }
+
+    return std::make_shared<CrossTensor>(std::move(out_tensor));
+}

native/extorch/src/lib.rs

@@ -114,6 +114,7 @@ rustler::init!(
         narrow_copy,
         nonzero,
         permute,
+        vstack,
 
         // Tensor comparing operations
         allclose,

native/extorch/src/native/tensor/ops.rs.in

@@ -152,3 +152,6 @@ fn permute(
     input: &SharedPtr<CrossTensor>,
     dims: Vec<i64>
 ) -> Result<SharedPtr<CrossTensor>>;
+
+/// Creates a new tensor by stacking vertically a sequence of tensors.
+fn vstack(tensors: TensorList, out: TensorOut) -> Result<SharedPtr<CrossTensor>>;

native/extorch/src/nifs/tensor_ops.rs

@@ -183,3 +183,10 @@ nif_impl!(
     input: TensorStruct<'a>,
     dims: Size
 );
+
+nif_impl!(
+    vstack,
+    TensorStruct<'a>,
+    tensors: TensorList,
+    out: TensorOut
+);

test/tensor/manipulation_test.exs

@@ -367,4 +367,19 @@ defmodule ExTorchTest.Tensor.ManipulationTest do
     out = ExTorch.permute(input, {2, -1, 0, 1})
     assert out.size == {5, 2, 3, 4}
   end
+
+  test "vstack/1" do
+    input = ExTorch.rand({5, 10, 2})
+    parts = [input[0..3], input[3..5]]
+    out = ExTorch.vstack(parts)
+    assert ExTorch.allclose(out, input)
+  end
+
+  test "vstack/2" do
+    input = ExTorch.rand({5, 10, 2})
+    parts = [input[0..3], input[3..5]]
+    out = ExTorch.empty_like(input)
+    ExTorch.vstack(parts, out)
+    assert ExTorch.allclose(out, input)
+  end
 end