utility function to flatten tensor

Author: ngc92

src/utilities/tensor.cpp

@@ -119,6 +119,14 @@ TensorShard shard_view(const Tensor& src, int idx, int num) {
     return TensorShard{shard, idx, num, src.Sizes};
 }
 
+Tensor flat_view(const Tensor& src) {
+    Tensor dst{src};
+    dst.Sizes.fill(0);
+    dst.Sizes[0] = src.nelem();
+    dst.Rank = 1;
+    return dst;
+}
+
 void visit(const std::function<void(Tensor&)>& func, SimpleTensorContainer& container) {
     auto cs = container.num_tensors();
     for(std::size_t i = 0; i < cs; ++i) {
@@ -168,6 +176,14 @@ GenericTensorContainer shard_empty_container(GenericTensorContainer&& c, int wor
     return std::move(c);
 }
 
+GenericTensorContainer flattened_view(const GenericTensorContainer& c) {
+    std::vector<Tensor> flats(c.num_tensors());
+    for (std::size_t i = 0; i < c.num_tensors(); ++i) {
+        flats.at(i) = flat_view(c.get_tensor(i));
+    }
+    return GenericTensorContainer{flats};
+}
+
 GenericTensorContainer shard_view(const GenericTensorContainer& c, int rank, int world) {
     std::vector<Tensor> shards(c.num_tensors());
     for (std::size_t i = 0; i < c.num_tensors(); ++i) {

src/utilities/tensor.h

@@ -160,4 +160,6 @@ class TensorShard : public Tensor {
 };
 
 TensorShard shard_view(const Tensor& src, int idx, int num);
+Tensor flat_view(const Tensor& src);
+
 #endif //LLMQ_SRC_UTILS_TENSOR_H

src/utilities/tensor_container.h

@@ -61,6 +61,9 @@ class GenericTensorContainer final : public SimpleTensorContainer {
 //! are `nullptr`, but sizes have been set up.
 GenericTensorContainer shard_empty_container(GenericTensorContainer&& c, int world);
 
+//! Flattens all tensors is the container.
+GenericTensorContainer flattened_view(const GenericTensorContainer& c);
+
 //! Shards a non-empty tensor container. The returned container's tensors are _views_ into
 //! the original container's tensors.
 GenericTensorContainer shard_view(const GenericTensorContainer& c, int rank, int world);