diff --git a/packages/react-native-executorch/cpp/core/tensor.cpp b/packages/react-native-executorch/cpp/core/tensor.cpp index 83bec4afaa..b424a87c49 100644 --- a/packages/react-native-executorch/cpp/core/tensor.cpp +++ b/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(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(optsObj.getProperty(rt, "offset").asNumber()); + } + + copyLen -= srcOffset; + + if (optsObj.hasProperty(rt, "length")) { + copyLen = static_cast(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)); }; diff --git a/packages/react-native-executorch/src/core/tensor.ts b/packages/react-native-executorch/src/core/tensor.ts index dcbddb6fc0..8f5c58e10c 100644 --- a/packages/react-native-executorch/src/core/tensor.ts +++ b/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.