[RNE Rewrite] feat: support offset and length options in Tensor.copyTo

packages/react-native-executorch/cpp/core/tensor.cpp

@@ -39,8 +39,8 @@ jsi::Value TensorHostObject::get(jsi::Runtime &rt, const jsi::PropNameID &name)
     if (nameStr == "copyTo") {
         auto self = shared_from_this();
         auto fnBody = [self](jsi::Runtime &rt, const jsi::Value & /*thisVal*/, const jsi::Value *args, size_t count) -> jsi::Value {
-            if (count != 1) {
-                throw jsi::JSError(rt, "copyTo: Usage: copyTo(dst)");
+            if (count != 1 && count != 2) {
+                throw jsi::JSError(rt, "copyTo: Usage: copyTo(dst, options?)");
             }
 
             if (!args[0].isObject() || !args[0].asObject(rt).isHostObject<TensorHostObject>(rt)) {
@@ -71,11 +71,31 @@ jsi::Value TensorHostObject::get(jsi::Runtime &rt, const jsi::PropNameID &name)
                 throw jsi::JSError(rt, "copyTo: dst tensor has been disposed");
             }
 
-            if (self->size_ != dst->size_) {
-                throw jsi::JSError(rt, "copyTo: size mismatch between src and dst tensors");
+            size_t srcOffset = 0;
+            size_t copyLen = self->numel_;
+            if (count == 2 && args[1].isObject()) {
+                auto optsObj = args[1].asObject(rt);
+                if (optsObj.hasProperty(rt, "offset")) {
+                    srcOffset = static_cast<size_t>(optsObj.getProperty(rt, "offset").asNumber());
+                }
+
+                copyLen -= srcOffset;
+
+                if (optsObj.hasProperty(rt, "length")) {
+                    copyLen = static_cast<size_t>(optsObj.getProperty(rt, "length").asNumber());
+                }
+            }
+
+            if (srcOffset + copyLen > self->numel_) {
+                throw jsi::JSError(rt, "copyTo: out of bounds offset and length for src tensor");
+            }
+
+            const auto elemSize = rnexecutorch::core::types::elementSize(self->dtype_);
+            if (copyLen * elemSize != dst->size_) {
+                throw jsi::JSError(rt, "copyTo: size mismatch between copy byte size and dst tensor size");
             }
 
-            std::memcpy(dst->data_.get(), self->data_.get(), self->size_);
+            std::memcpy(dst->data_.get(), self->data_.get() + (srcOffset * elemSize), copyLen * elemSize);
 
             return jsi::Value(rt, args[0].asObject(rt));
         };

packages/react-native-executorch/src/core/tensor.ts

@@ -28,11 +28,17 @@ export type Tensor = {
   readonly numel: number;
 
   /**
-   * Copies this tensor's data into another tensor with equal number of bytes.
+   * Copies this tensor's data into another tensor.
    * @param dst The destination tensor to copy data into.
+   * @param options Optional configuration for the copy operation.
+   * @param options.offset The start offset in elements in the source tensor.
+   * Defaults to `0`.
+   * @param options.length The number of elements to copy. Defaults to
+   * `numel - offset`, i.e. copies from `offset` to the end of the source
+   * tensor.
    * @returns The destination tensor `dst`.
    */
-  copyTo(dst: Tensor): Tensor;
+  copyTo(dst: Tensor, options?: { offset?: number; length?: number }): Tensor;
 
   /**
    * Releases the underlying native C++ memory held by this tensor.