Allow data casting along with cloning. (#15511)

Summary:

.

Differential Revision: D86070965

Author: shoumikhin

extension/apple/ExecuTorch/Exported/ExecuTorch+Tensor.swift

@@ -741,6 +741,14 @@ public final class Tensor<T: Scalar>: Equatable {
     Tensor<T>(anyTensor.copy())
   }
 
+  /// Returns a copy of the tensor, converted to the specified scalar type.
+  ///
+  /// - Parameter dataType: The target scalar type.
+  /// - Returns: A new tensor with the same shape and metadata but converted elements.
+  public func copy<U: Scalar>(to dataType: U.Type) -> Tensor<U> {
+    Tensor<U>(anyTensor.copy(to: U.dataType))
+  }
+
   /// Calls the closure with a typed, immutable buffer pointer over the tensor’s elements.
   ///
   /// - Parameter body: A closure that receives an `UnsafeBufferPointer<T>` bound to the tensor’s data.

extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h

@@ -168,6 +168,16 @@ __attribute__((objc_subclassing_restricted))
  */
 - (instancetype)copy;
 
+/**
+ * Creates a deep copy of the tensor, potentially casting to a new data type.
+ * The new tensor will have its own copy of the data.
+ *
+ * @param dataType The desired data type for the new tensor.
+ * @return A new ExecuTorchTensor instance that is a duplicate (and possibly casted) of the current tensor.
+*/
+- (instancetype)copyToDataType:(ExecuTorchDataType)dataType
+    NS_SWIFT_NAME(copy(to:));
+
 /**
  * Executes a block with a pointer to the tensor's immutable byte data.
  *

extension/apple/ExecuTorch/Exported/ExecuTorchTensor.mm

@@ -147,6 +147,11 @@ - (instancetype)copyWithZone:(nullable NSZone *)zone {
   return [[ExecuTorchTensor allocWithZone:zone] initWithNativeInstance:&tensor];
 }
 
+- (instancetype)copyToDataType:(ExecuTorchDataType)dataType {
+  auto tensor = clone_tensor_ptr(_tensor, static_cast<ScalarType>(dataType));
+  return [[ExecuTorchTensor alloc] initWithNativeInstance:&tensor];
+}
+
 - (void *)nativeInstance {
   return &_tensor;
 }

extension/apple/ExecuTorch/__tests__/TensorTest.swift

@@ -182,6 +182,43 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor1.count, tensor2.count)
   }
 
+  func testCopyToSameDataType() {
+    let tensor1 = Tensor<Float>([1, 2, 3, 4], shape: [2, 2])
+    let tensor2 = tensor1.copy(to: Float.self)
+    XCTAssertEqual(tensor2.dataType, .float)
+    XCTAssertEqual(tensor2.shape, [2, 2])
+    XCTAssertEqual(tensor2.strides, tensor1.strides)
+    XCTAssertEqual(tensor2.dimensionOrder, tensor1.dimensionOrder)
+    XCTAssertEqual(tensor2.scalars(), [1, 2, 3, 4])
+  }
+
+  func testCopyToDifferentDataTypeKeepsSourceAlive() {
+    var data = [10.0, 20.0, 30.0, 40.0]
+    let tensor1 = data.withUnsafeMutableBytes {
+      Tensor<Double>(bytesNoCopy: $0.baseAddress!, shape: [2, 2])
+    }
+    let tensor2 = tensor1.copy(to: Float.self)
+    data[0] = 999.0
+    XCTAssertEqual(tensor2.dataType, .float)
+    XCTAssertEqual(tensor2.shape, [2, 2])
+    XCTAssertEqual(tensor2.scalars(), [10.0, 20.0, 30.0, 40.0])
+  }
+
+  func testCopyToPreservesShapeAndOrderOn2D() {
+    let tensor1 = Tensor<Int32>(
+      [1, 2, 3, 4, 5, 6],
+      shape: [2, 3],
+      strides: [3, 1],
+      dimensionOrder: [0, 1]
+    )
+    let tensor2 = tensor1.copy(to: Double.self)
+    XCTAssertEqual(tensor2.shape, [2, 3])
+    XCTAssertEqual(tensor2.strides, [3, 1])
+    XCTAssertEqual(tensor2.dimensionOrder, [0, 1])
+    XCTAssertEqual(tensor2.count, 6)
+    XCTAssertEqual(tensor2.scalars(), [1, 2, 3, 4, 5, 6])
+  }
+
   func testResize() {
     var data: [Int] = [1, 2, 3, 4]
     let tensor = data.withUnsafeMutableBytes {