Add making JsTensor from iterator

Author: Conarnar

extension/wasm/test/unittests.js

@@ -36,6 +36,16 @@ describe("Tensor", () => {
         tensor.delete();
     });
 
+    test("fromGenerator", () => {
+        function* generator() {
+            yield* [1, 2, 3, 4];
+        }
+        const tensor = et.Tensor.fromIter([2, 2], generator());
+        expect(tensor.data).toEqual(new Float32Array([1, 2, 3, 4]));
+        expect(tensor.sizes).toEqual([2, 2]);
+        tensor.delete();
+    });
+
     test("fromArray wrong size", () => {
         expect(() => et.Tensor.fromArray([3, 2], [1, 2, 3, 4])).toThrow();
     });

extension/wasm/wasm_bindings.cpp

@@ -83,6 +83,19 @@ inline void assert_valid_numel(
       data.size());
 }
 
+template <typename T>
+std::vector<T> convertJSGeneratorToNumberVector(val generator) {
+  std::vector<T> data;
+  while (true) {
+    val next = generator.call<val>("next");
+    if (next["done"].as<bool>()) {
+      break;
+    }
+    data.push_back(next["value"].as<T>());
+  }
+  return data;
+}
+
 class JsTensor {
  public:
   JsTensor() = delete;
@@ -204,6 +217,49 @@ class JsTensor {
     }
   }
 
+  static std::unique_ptr<JsTensor> from_iter(
+      val_array<int> sizes,
+      val_array<val> data,
+      val type = val::undefined(),
+      val_array<int> dim_order = val::undefined(),
+      val_array<int> strides = val::undefined()) {
+    auto sizes_vec =
+        convertJSArrayToNumberVector<executorch::aten::SizesType>(sizes);
+
+    auto dim_order_vec = dim_order.isUndefined()
+        ? std::vector<executorch::aten::DimOrderType>()
+        : convertJSArrayToNumberVector<executorch::aten::DimOrderType>(
+              dim_order);
+    auto strides_vec = strides.isUndefined()
+        ? std::vector<executorch::aten::StridesType>()
+        : convertJSArrayToNumberVector<executorch::aten::StridesType>(strides);
+
+    // If type is undefined, infer the type from the data.
+    // Assume it is a Bigint if not Number.
+    ScalarType scalar_type = type.isUndefined()
+        ? (data["length"].as<size_t>() == 0 || data[0].isNumber()
+               ? ScalarType::Float
+               : ScalarType::Long)
+        : type.as<ScalarType>();
+    switch (scalar_type) {
+#define JS_CASE_FROM_ITER_VECTOR_TYPE(T, NAME)                 \
+  case ScalarType::NAME: {                                     \
+    auto data_vec = convertJSGeneratorToNumberVector<T>(data); \
+    assert_valid_numel(data_vec, sizes_vec);                   \
+    TensorPtr tensor = make_tensor_ptr(                        \
+        std::move(sizes_vec),                                  \
+        std::move(data_vec),                                   \
+        std::move(dim_order_vec),                              \
+        std::move(strides_vec),                                \
+        ScalarType::NAME);                                     \
+    return std::make_unique<JsTensor>(std::move(tensor));      \
+  }
+      JS_FORALL_SUPPORTED_TENSOR_TYPES(JS_CASE_FROM_ITER_VECTOR_TYPE)
+      default:
+        THROW_JS_ERROR(TypeError, "Unsupported Tensor type: %d", scalar_type);
+    }
+  }
+
  private:
   TensorPtr tensor_;
 };
@@ -448,6 +504,7 @@ EMSCRIPTEN_BINDINGS(WasmBindings) {
       .class_function("ones", &JsTensor::ones)
       .class_function("full", &JsTensor::full)
       .class_function("fromArray", &JsTensor::from_array)
+      .class_function("fromIter", &JsTensor::from_iter)
       .property("scalarType", &JsTensor::get_scalar_type)
       .property("data", &JsTensor::get_data)
       .property("sizes", &JsTensor::get_sizes);