diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorch+Tensor.swift b/extension/apple/ExecuTorch/Exported/ExecuTorch+Tensor.swift
index 29af8f78a5a..b325000ed23 100644
--- a/extension/apple/ExecuTorch/Exported/ExecuTorch+Tensor.swift
+++ b/extension/apple/ExecuTorch/Exported/ExecuTorch+Tensor.swift
@@ -54,7 +54,7 @@ public extension Tensor {
   func withUnsafeBytes<T: Scalar, R>(_ body: (UnsafeBufferPointer<T>) throws -> R) throws -> R {
     guard dataType == T.dataType else { throw Error(code: .invalidArgument) }
     var result: Result<R, Error>?
-    bytes { pointer, count, _ in
+    __bytes { pointer, count, _ in
       result = Result { try body(
         UnsafeBufferPointer(
           start: pointer.assumingMemoryBound(to: T.self),
@@ -74,7 +74,7 @@ public extension Tensor {
   func withUnsafeMutableBytes<T: Scalar, R>(_ body: (UnsafeMutableBufferPointer<T>) throws -> R) throws -> R {
     guard dataType == T.dataType else { throw Error(code: .invalidArgument) }
     var result: Result<R, Error>?
-    mutableBytes { pointer, count, _ in
+    __mutableBytes { pointer, count, _ in
       result = Result { try body(
         UnsafeMutableBufferPointer(
           start: pointer.assumingMemoryBound(to: T.self),
diff --git a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h
index e832845d6ba..5b130da56c9 100644
--- a/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h
+++ b/extension/apple/ExecuTorch/Exported/ExecuTorchTensor.h
@@ -176,7 +176,7 @@ __attribute__((deprecated("This API is experimental.")))
  *   - and the data type.
  */
 - (void)bytesWithHandler:(NS_NOESCAPE void (^)(const void *pointer, NSInteger count, ExecuTorchDataType dataType))handler
-    NS_SWIFT_NAME(bytes(_:));
+    NS_REFINED_FOR_SWIFT;
 
 /**
  * Executes a block with a pointer to the tensor's mutable byte data.
@@ -187,7 +187,7 @@ __attribute__((deprecated("This API is experimental.")))
  *   - and the data type.
  */
 - (void)mutableBytesWithHandler:(NS_NOESCAPE void (^)(void *pointer, NSInteger count, ExecuTorchDataType dataType))handler
-    NS_SWIFT_NAME(mutableBytes(_:));
+    NS_REFINED_FOR_SWIFT;
 
 /**
  * Resizes the tensor to a new shape.
diff --git a/extension/apple/ExecuTorch/__tests__/TensorTest.swift b/extension/apple/ExecuTorch/__tests__/TensorTest.swift
index 689a514403f..052b84ae5f8 100644
--- a/extension/apple/ExecuTorch/__tests__/TensorTest.swift
+++ b/extension/apple/ExecuTorch/__tests__/TensorTest.swift
@@ -68,13 +68,9 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.dimensionOrder, [0, 1])
     XCTAssertEqual(tensor.shapeDynamism, .dynamicBound)
     XCTAssertEqual(tensor.count, 6)
-
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(dataType, .float)
-      XCTAssertEqual(count, 6)
-      XCTAssertEqual(size(ofDataType: dataType), 4)
-      XCTAssertEqual(Array(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Float.self), count: count)), data)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { buffer in
+      XCTAssertEqual(Array(buffer), data)
+    })
   }
 
   func testInitBytes() {
@@ -91,13 +87,9 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.dimensionOrder, [0, 1])
     XCTAssertEqual(tensor.shapeDynamism, .dynamicBound)
     XCTAssertEqual(tensor.count, 6)
-
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(dataType, .double)
-      XCTAssertEqual(count, 6)
-      XCTAssertEqual(size(ofDataType: dataType), 8)
-      XCTAssertEqual(Array(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Double.self), count: count)).map { $0 + 1 }, data)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { buffer in
+      XCTAssertEqual(Array(buffer).map { $0 + 1 }, data)
+    })
   }
 
   func testInitData() {
@@ -105,9 +97,9 @@ class TensorTest: XCTestCase {
     let data = Data(bytes: dataArray, count: dataArray.count * MemoryLayout<Float>.size)
     let tensor = Tensor(data: data, shape: [4], dataType: .float)
     XCTAssertEqual(tensor.count, 4)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(Array(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Float.self), count: count)), dataArray)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { buffer in
+      XCTAssertEqual(Array(buffer), dataArray)
+    })
   }
 
   func testWithCustomStridesAndDimensionOrder() {
@@ -123,10 +115,9 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.strides, [1, 2])
     XCTAssertEqual(tensor.dimensionOrder, [1, 0])
     XCTAssertEqual(tensor.count, 4)
-
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(Array(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Float.self), count: count)), data)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { buffer in
+      XCTAssertEqual(Array(buffer), data)
+    })
   }
 
   func testMutableBytes() {
@@ -134,41 +125,14 @@ class TensorTest: XCTestCase {
     let tensor = data.withUnsafeMutableBytes {
       Tensor(bytes: $0.baseAddress!, shape: [4], dataType: .int)
     }
-    tensor.mutableBytes { pointer, count, dataType in
-      XCTAssertEqual(dataType, .int)
-      let buffer = pointer.assumingMemoryBound(to: Int32.self)
-      for i in 0..<count {
-        buffer[i] *= 2
-      }
-    }
-    tensor.bytes { pointer, count, dataType in
-      let updatedData = Array(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Int32.self), count: count))
-      XCTAssertEqual(updatedData, [2, 4, 6, 8])
-    }
-  }
-
-  func testWithUnsafeBytes() throws {
-    var data: [Float] = [1, 2, 3, 4, 5, 6]
-    let tensor = data.withUnsafeMutableBytes {
-      Tensor(bytesNoCopy: $0.baseAddress!, shape: [2, 3], dataType: .float)
-    }
-    let array = try tensor.withUnsafeBytes([Float].init)
-    XCTAssertEqual(array, data)
-  }
-
-  func testWithUnsafeMutableBytes() throws {
-    var data = [1, 2, 3, 4]
-    let tensor = data.withUnsafeMutableBytes {
-      Tensor(bytes: $0.baseAddress!, shape: [4], dataType: .long)
-    }
-    try tensor.withUnsafeMutableBytes { (buffer: UnsafeMutableBufferPointer<Int>) in
+    XCTAssertNoThrow(try tensor.withUnsafeMutableBytes { (buffer: UnsafeMutableBufferPointer<Int32>) in
       for i in buffer.indices {
         buffer[i] *= 2
       }
-    }
-    try tensor.withUnsafeBytes { buffer in
+    })
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { (buffer: UnsafeBufferPointer<Int32>) in
       XCTAssertEqual(Array(buffer), [2, 4, 6, 8])
-    }
+    })
   }
 
   func testInitWithTensor() {
@@ -202,18 +166,17 @@ class TensorTest: XCTestCase {
   func testResize() {
     var data: [Int] = [1, 2, 3, 4]
     let tensor = data.withUnsafeMutableBytes {
-      Tensor(bytesNoCopy: $0.baseAddress!, shape: [4, 1], dataType: .int)
+      Tensor(bytesNoCopy: $0.baseAddress!, shape: [4, 1], dataType: .long)
     }
     XCTAssertNoThrow(try tensor.resize(to: [2, 2]))
-    XCTAssertEqual(tensor.dataType, .int)
+    XCTAssertEqual(tensor.dataType, .long)
     XCTAssertEqual(tensor.shape, [2, 2])
     XCTAssertEqual(tensor.strides, [2, 1])
     XCTAssertEqual(tensor.dimensionOrder, [0, 1])
     XCTAssertEqual(tensor.count, 4)
-
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(Array(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Int.self), count: count)), data)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { buffer in
+      XCTAssertEqual(Array(buffer), data)
+    })
   }
 
   func testResizeError() {
@@ -255,9 +218,9 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.strides, [1])
     XCTAssertEqual(tensor.dimensionOrder, [0])
     XCTAssertEqual(tensor.count, 6)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(Array(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: UInt8.self), count: count)), data)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { buffer in
+      XCTAssertEqual(Array(buffer), data)
+    })
   }
 
   func testInitScalarsInt8() {
@@ -268,9 +231,9 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.strides, [1])
     XCTAssertEqual(tensor.dimensionOrder, [0])
     XCTAssertEqual(tensor.count, 6)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(Array(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Int8.self), count: count)), data)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { buffer in
+      XCTAssertEqual(Array(buffer), data)
+    })
   }
 
   func testInitScalarsInt16() {
@@ -281,9 +244,9 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.strides, [1])
     XCTAssertEqual(tensor.dimensionOrder, [0])
     XCTAssertEqual(tensor.count, 6)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(Array(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Int16.self), count: count)), data)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { buffer in
+      XCTAssertEqual(Array(buffer), data)
+    })
   }
 
   func testInitScalarsInt32() {
@@ -294,9 +257,9 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.strides, [1])
     XCTAssertEqual(tensor.dimensionOrder, [0])
     XCTAssertEqual(tensor.count, 6)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(Array(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Int32.self), count: count)), data)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { buffer in
+      XCTAssertEqual(Array(buffer), data)
+    })
   }
 
   func testInitScalarsInt64() {
@@ -307,9 +270,9 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.strides, [1])
     XCTAssertEqual(tensor.dimensionOrder, [0])
     XCTAssertEqual(tensor.count, 6)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(Array(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Int64.self), count: count)), data)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { buffer in
+      XCTAssertEqual(Array(buffer), data)
+    })
   }
 
   func testInitScalarsFloat() {
@@ -320,9 +283,9 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.strides, [1])
     XCTAssertEqual(tensor.dimensionOrder, [0])
     XCTAssertEqual(tensor.count, 6)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(Array(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Float.self), count: count)), data)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { buffer in
+      XCTAssertEqual(Array(buffer), data)
+    })
   }
 
   func testInitScalarsDouble() {
@@ -333,9 +296,9 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.strides, [1])
     XCTAssertEqual(tensor.dimensionOrder, [0])
     XCTAssertEqual(tensor.count, 6)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(Array(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Double.self), count: count)), data)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { buffer in
+      XCTAssertEqual(Array(buffer), data)
+    })
   }
 
   func testInitScalarsBool() {
@@ -346,9 +309,9 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.strides, [1])
     XCTAssertEqual(tensor.dimensionOrder, [0])
     XCTAssertEqual(tensor.count, 6)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(Array(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Bool.self), count: count)), data)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { buffer in
+      XCTAssertEqual(Array(buffer), data)
+    })
   }
 
   func testInitScalarsUInt16() {
@@ -359,9 +322,9 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.strides, [1])
     XCTAssertEqual(tensor.dimensionOrder, [0])
     XCTAssertEqual(tensor.count, 6)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(Array(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: UInt16.self), count: count)), data)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { buffer in
+      XCTAssertEqual(Array(buffer), data)
+    })
   }
 
   func testInitScalarsUInt32() {
@@ -372,9 +335,9 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.strides, [1])
     XCTAssertEqual(tensor.dimensionOrder, [0])
     XCTAssertEqual(tensor.count, 6)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(Array(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: UInt32.self), count: count)), data)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { buffer in
+      XCTAssertEqual(Array(buffer), data)
+    })
   }
 
   func testInitScalarsUInt64() {
@@ -385,9 +348,9 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.strides, [1])
     XCTAssertEqual(tensor.dimensionOrder, [0])
     XCTAssertEqual(tensor.count, 6)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(Array(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: UInt64.self), count: count)), data)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { buffer in
+      XCTAssertEqual(Array(buffer), data)
+    })
   }
 
   func testInitScalarsInt() {
@@ -398,9 +361,9 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.strides, [1])
     XCTAssertEqual(tensor.dimensionOrder, [0])
     XCTAssertEqual(tensor.count, 6)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(Array(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Int.self), count: count)), data)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { buffer in
+      XCTAssertEqual(Array(buffer), data)
+    })
   }
 
   func testInitScalarsUInt() {
@@ -411,9 +374,9 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.strides, [1])
     XCTAssertEqual(tensor.dimensionOrder, [0])
     XCTAssertEqual(tensor.count, 6)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(Array(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: UInt.self), count: count)), data)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { buffer in
+      XCTAssertEqual(Array(buffer), data)
+    })
   }
 
   func testInitInt8() {
@@ -423,9 +386,9 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.strides, [])
     XCTAssertEqual(tensor.dimensionOrder, [])
     XCTAssertEqual(tensor.count, 1)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Int8.self), count: count).first, 42)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { (buffer: UnsafeBufferPointer<Int8>) in
+      XCTAssertEqual(Array(buffer).first, 42)
+    })
   }
 
   func testInitInt16() {
@@ -435,9 +398,9 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.strides, [])
     XCTAssertEqual(tensor.dimensionOrder, [])
     XCTAssertEqual(tensor.count, 1)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Int16.self), count: count).first, 42)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { (buffer: UnsafeBufferPointer<Int16>) in
+      XCTAssertEqual(Array(buffer).first, 42)
+    })
   }
 
   func testInitInt32() {
@@ -447,9 +410,9 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.strides, [])
     XCTAssertEqual(tensor.dimensionOrder, [])
     XCTAssertEqual(tensor.count, 1)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Int32.self), count: count).first, 42)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { (buffer: UnsafeBufferPointer<Int32>) in
+      XCTAssertEqual(Array(buffer).first, 42)
+    })
   }
 
   func testInitInt64() {
@@ -459,9 +422,9 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.strides, [])
     XCTAssertEqual(tensor.dimensionOrder, [])
     XCTAssertEqual(tensor.count, 1)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Int64.self), count: count).first, 42)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { (buffer: UnsafeBufferPointer<Int64>) in
+      XCTAssertEqual(Array(buffer).first, 42)
+    })
   }
 
   func testInitUInt8() {
@@ -471,9 +434,9 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.strides, [])
     XCTAssertEqual(tensor.dimensionOrder, [])
     XCTAssertEqual(tensor.count, 1)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: UInt8.self), count: count).first, 42)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { (buffer: UnsafeBufferPointer<UInt8>) in
+      XCTAssertEqual(Array(buffer).first, 42)
+    })
   }
 
   func testInitUInt16() {
@@ -483,9 +446,9 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.strides, [])
     XCTAssertEqual(tensor.dimensionOrder, [])
     XCTAssertEqual(tensor.count, 1)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: UInt16.self), count: count).first, 42)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { (buffer: UnsafeBufferPointer<UInt16>) in
+      XCTAssertEqual(Array(buffer).first, 42)
+    })
   }
 
   func testInitUInt32() {
@@ -495,9 +458,9 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.strides, [])
     XCTAssertEqual(tensor.dimensionOrder, [])
     XCTAssertEqual(tensor.count, 1)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: UInt32.self), count: count).first, 42)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { (buffer: UnsafeBufferPointer<UInt32>) in
+      XCTAssertEqual(Array(buffer).first, 42)
+    })
   }
 
   func testInitUInt64() {
@@ -507,9 +470,9 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.strides, [])
     XCTAssertEqual(tensor.dimensionOrder, [])
     XCTAssertEqual(tensor.count, 1)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: UInt64.self), count: count).first, 42)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { (buffer: UnsafeBufferPointer<UInt64>) in
+      XCTAssertEqual(Array(buffer).first, 42)
+    })
   }
 
   func testInitBool() {
@@ -519,9 +482,9 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.strides, [])
     XCTAssertEqual(tensor.dimensionOrder, [])
     XCTAssertEqual(tensor.count, 1)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Bool.self), count: count).first, true)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { (buffer: UnsafeBufferPointer<Bool>) in
+      XCTAssertEqual(Array(buffer).first, true)
+    })
   }
 
   func testInitFloat() {
@@ -531,9 +494,9 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.strides, [])
     XCTAssertEqual(tensor.dimensionOrder, [])
     XCTAssertEqual(tensor.count, 1)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Float.self), count: count).first, 42.0)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { (buffer: UnsafeBufferPointer<Float>) in
+      XCTAssertEqual(Array(buffer).first, 42)
+    })
   }
 
   func testInitDouble() {
@@ -543,9 +506,9 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.strides, [])
     XCTAssertEqual(tensor.dimensionOrder, [])
     XCTAssertEqual(tensor.count, 1)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Double.self), count: count).first, 42.0)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { (buffer: UnsafeBufferPointer<Double>) in
+      XCTAssertEqual(Array(buffer).first, 42.0)
+    })
   }
 
   func testInitInt() {
@@ -555,9 +518,9 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.strides, [])
     XCTAssertEqual(tensor.dimensionOrder, [])
     XCTAssertEqual(tensor.count, 1)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Int.self), count: count).first, 42)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { buffer in
+      XCTAssertEqual(Array(buffer).first, 42)
+    })
   }
 
   func testInitUInt() {
@@ -567,20 +530,20 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.strides, [])
     XCTAssertEqual(tensor.dimensionOrder, [])
     XCTAssertEqual(tensor.count, 1)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: UInt.self), count: count).first, 42)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { (buffer: UnsafeBufferPointer<UInt64>) in
+      XCTAssertEqual(Array(buffer).first, 42)
+    })
   }
 
   func testEmpty() {
     let tensor = Tensor.empty(shape: [3, 4], dataType: .float)
     XCTAssertEqual(tensor.shape, [3, 4])
     XCTAssertEqual(tensor.count, 12)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertNotNil(pointer)
-      XCTAssertEqual(count, 12)
-      XCTAssertEqual(dataType, .float)
-    }
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { (buffer: UnsafeBufferPointer<Float>) in
+      XCTAssertNotNil(buffer.baseAddress)
+      XCTAssertEqual(buffer.count, 12)
+      XCTAssertEqual(tensor.dataType, .float)
+    })
   }
 
   func testEmptyLike() {
@@ -596,87 +559,76 @@ class TensorTest: XCTestCase {
     let tensor = Tensor.full(shape: [2, 2], scalar: 7, dataType: .int)
     XCTAssertEqual(tensor.shape, [2, 2])
     XCTAssertEqual(tensor.count, 4)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(dataType, .int)
-      let buffer = UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Int32.self), count: count)
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { (buffer: UnsafeBufferPointer<Int32>) in
       for value in buffer {
         XCTAssertEqual(value, 7)
       }
-    }
+    })
   }
 
   func testFullLike() {
     let other = Tensor.empty(shape: [2, 2], dataType: .int)
     let tensor = Tensor.full(like: other, scalar: 42, dataType: .float)
     XCTAssertEqual(tensor.shape, other.shape)
-    tensor.bytes { pointer, count, dataType in
-      let buffer = UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Float.self), count: count)
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { (buffer: UnsafeBufferPointer<Float>) in
       for value in buffer {
         XCTAssertEqual(value, 42.0)
       }
-    }
+    })
   }
 
   func testOnes() {
     let tensor = Tensor.ones(shape: [2, 3], dataType: .float)
     XCTAssertEqual(tensor.shape, [2, 3])
     XCTAssertEqual(tensor.count, 6)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(dataType, .float)
-      let buffer = UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Float.self), count: count)
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { (buffer: UnsafeBufferPointer<Float>) in
       for value in buffer {
         XCTAssertEqual(value, 1.0)
       }
-    }
+    })
   }
 
   func testOnesLike() {
     let other = Tensor.empty(shape: [2, 4], dataType: .double)
     let tensor = Tensor.ones(like: other)
     XCTAssertEqual(tensor.shape, other.shape)
-    tensor.bytes { pointer, count, dataType in
-      let buffer = UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Double.self), count: count)
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { (buffer: UnsafeBufferPointer<Double>) in
       for value in buffer {
         XCTAssertEqual(value, 1.0)
       }
-    }
+    })
   }
 
   func testZeros() {
     let tensor = Tensor.zeros(shape: [2, 3], dataType: .double)
     XCTAssertEqual(tensor.shape, [2, 3])
     XCTAssertEqual(tensor.count, 6)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(dataType, .double)
-      let buffer = UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Double.self), count: count)
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { (buffer: UnsafeBufferPointer<Double>) in
       for value in buffer {
         XCTAssertEqual(value, 0)
       }
-    }
+    })
   }
 
   func testZerosLike() {
     let other = Tensor.full(shape: [3, 2], scalar: 9, dataType: .int)
     let tensor = Tensor.zeros(like: other)
     XCTAssertEqual(tensor.shape, other.shape)
-    tensor.bytes { pointer, count, dataType in
-      let buffer = UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Int32.self), count: count)
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { (buffer: UnsafeBufferPointer<Int32>) in
       for value in buffer {
         XCTAssertEqual(value, 0)
       }
-    }
+    })
   }
 
   func testRandom() {
     let tensor = Tensor.rand(shape: [3, 3], dataType: .float)
     XCTAssertEqual(tensor.shape, [3, 3])
     XCTAssertEqual(tensor.count, 9)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(dataType, .float)
-      let buffer = UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Float.self), count: count)
-      let uniqueValues = Set(buffer.map { $0 })
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { (buffer: UnsafeBufferPointer<Float>) in
+      let uniqueValues = Set(buffer)
       XCTAssertTrue(uniqueValues.count > 1)
-    }
+    })
   }
 
   func testRandomLike() {
@@ -686,15 +638,13 @@ class TensorTest: XCTestCase {
     XCTAssertEqual(tensor.count, other.count)
   }
 
- func testRandomNormal() {
+  func testRandomNormal() {
     let tensor = Tensor.randn(shape: [4], dataType: .double)
     XCTAssertEqual(tensor.shape, [4])
     XCTAssertEqual(tensor.count, 4)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(dataType, .double)
-      let buffer = UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Double.self), count: count)
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { (buffer: UnsafeBufferPointer<Double>) in
       XCTAssertEqual(buffer.count, 4)
-    }
+    })
   }
 
   func testRandomNormalLike() {
@@ -708,23 +658,20 @@ class TensorTest: XCTestCase {
     let tensor = Tensor.randint(low: 10, high: 20, shape: [5], dataType: .int)
     XCTAssertEqual(tensor.shape, [5])
     XCTAssertEqual(tensor.count, 5)
-    tensor.bytes { pointer, count, dataType in
-      XCTAssertEqual(dataType, .int)
-      let buffer = UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Int32.self), count: count)
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { (buffer: UnsafeBufferPointer<Int32>) in
       for value in buffer {
         XCTAssertTrue(value >= 10 && value < 20)
       }
-    }
+    })
   }
 
   func testRandomIntegerLike() {
     let other = Tensor.ones(shape: [5], dataType: .int)
     let tensor = Tensor.randint(like: other, low: 100, high: 200)
-    tensor.bytes { pointer, count, dataType in
-      let buffer = UnsafeBufferPointer(start: pointer.assumingMemoryBound(to: Int32.self), count: count)
+    XCTAssertNoThrow(try tensor.withUnsafeBytes { (buffer: UnsafeBufferPointer<Int32>) in
       for value in buffer {
         XCTAssertTrue(value >= 100 && value < 200)
       }
-    }
+    })
   }
 }