ztensor 0.1.0__tar.gz → 0.1.2__tar.gz

{ztensor-0.1.0 → ztensor-0.1.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ztensor
-Version: 0.1.0
+Version: 0.1.2
 Classifier: Programming Language :: Rust
 Classifier: Programming Language :: Python :: 3
 Classifier: License :: OSI Approved :: MIT License
@@ -9,6 +9,7 @@ Classifier: Intended Audience :: Developers
 Classifier: Topic :: Scientific/Engineering
 Requires-Dist: numpy
 Requires-Dist: cffi
+Requires-Dist: ml-dtypes
 License-File: LICENSE
 Summary: Python bindings for the zTensor library.
 Author: In Gim <in.gim@yale.edu>

{ztensor-0.1.0 → ztensor-0.1.2}/pyproject.toml

@@ -10,7 +10,7 @@ build-backend = "maturin"
 [project]
 # The name of your package on PyPI.
 name = "ztensor"
-version = "0.1.0"
+version = "0.1.2"
 description = "Python bindings for the zTensor library."
 readme = "README.md" # It's good practice to have a README.
 authors = [
@@ -27,7 +27,8 @@ classifiers = [
 ]
 dependencies = [
     "numpy",
-    "cffi"
+    "cffi",
+    "ml_dtypes"
 ]
 
 [project.urls]

{ztensor-0.1.0 → ztensor-0.1.2}/python/examples/basic.py

@@ -2,9 +2,17 @@ import os
 
 import numpy as np
 from ztensor import Writer, Reader, ZTensorError
+import torch
 
+model_path = "llama1b.zt"
 
-file_path = "test_tensors.zt"
+with Reader(model_path) as reader:
+    tensor_data_np = reader.read_tensor("model.layers.1.self_attn.k_proj.weight", to="torch")
+    print(tensor_data_np.dtype)
+
+
+
+file_path = "../test_tensors.zt"
 
 # --- Write Tensors ---
 print(f"--- Writing to {file_path} ---")

{ztensor-0.1.0 → ztensor-0.1.2}/python/ztensor/__init__.py

@@ -1,6 +1,23 @@
 import numpy as np
 from .ztensor import ffi, lib
 
+# --- Optional PyTorch Import ---
+try:
+    import torch
+
+    TORCH_AVAILABLE = True
+except ImportError:
+    TORCH_AVAILABLE = False
+
+# --- Optional ml_dtypes for bfloat16 in NumPy ---
+try:
+    from ml_dtypes import bfloat16 as np_bfloat16
+
+    ML_DTYPES_AVAILABLE = True
+except ImportError:
+    np_bfloat16 = None
+    ML_DTYPES_AVAILABLE = False
+
 
 # --- Pythonic Wrapper ---
 class ZTensorError(Exception):
@@ -11,7 +28,6 @@ class ZTensorError(Exception):
 # A custom ndarray subclass to safely manage the lifetime of the CFFI pointer.
 class _ZTensorView(np.ndarray):
     def __new__(cls, buffer, dtype, shape, view_ptr):
-        # Create an array from the buffer, reshape it, and cast it to our custom type.
         obj = np.frombuffer(buffer, dtype=dtype).reshape(shape).view(cls)
         # Attach the object that owns the memory to an attribute.
         obj._owner = view_ptr
@@ -44,23 +60,36 @@ def _check_status(status, func_name=""):
         raise ZTensorError(f"Error in {func_name}: {_get_last_error()}")
 
 
-# Type Mappings between NumPy and ztensor
+# --- Type Mappings ---
+# NumPy Mappings
 DTYPE_NP_TO_ZT = {
-    np.dtype('float64'): 'float64', np.dtype('float32'): 'float32',
+    np.dtype('float64'): 'float64', np.dtype('float32'): 'float32', np.dtype('float16'): 'float16',
     np.dtype('int64'): 'int64', np.dtype('int32'): 'int32',
     np.dtype('int16'): 'int16', np.dtype('int8'): 'int8',
     np.dtype('uint64'): 'uint64', np.dtype('uint32'): 'uint32',
     np.dtype('uint16'): 'uint16', np.dtype('uint8'): 'uint8',
     np.dtype('bool'): 'bool',
 }
+if ML_DTYPES_AVAILABLE:
+    DTYPE_NP_TO_ZT[np.dtype(np_bfloat16)] = 'bfloat16'
 DTYPE_ZT_TO_NP = {v: k for k, v in DTYPE_NP_TO_ZT.items()}
 
+# PyTorch Mappings (if available)
+if TORCH_AVAILABLE:
+    DTYPE_TORCH_TO_ZT = {
+        torch.float64: 'float64', torch.float32: 'float32', torch.float16: 'float16',
+        torch.bfloat16: 'bfloat16',
+        torch.int64: 'int64', torch.int32: 'int32',
+        torch.int16: 'int16', torch.int8: 'int8',
+        torch.uint8: 'uint8', torch.bool: 'bool',
+    }
+    DTYPE_ZT_TO_TORCH = {v: k for k, v in DTYPE_TORCH_TO_ZT.items()}
+
 
 class TensorMetadata:
     """A Pythonic wrapper around the CTensorMetadata pointer."""
 
     def __init__(self, meta_ptr):
-        # The pointer is now automatically garbage collected by CFFI when this object dies.
         self._ptr = ffi.gc(meta_ptr, lib.ztensor_metadata_free)
         _check_ptr(self._ptr, "TensorMetadata constructor")
         self._name = None
@@ -72,7 +101,6 @@ class TensorMetadata:
         if self._name is None:
             name_ptr = lib.ztensor_metadata_get_name(self._ptr)
             _check_ptr(name_ptr, "get_name")
-            # ffi.string creates a copy, so we must free the Rust-allocated original.
             self._name = ffi.string(name_ptr).decode('utf-8')
             lib.ztensor_free_string(name_ptr)
         return self._name
@@ -82,7 +110,6 @@ class TensorMetadata:
         if self._dtype_str is None:
             dtype_ptr = lib.ztensor_metadata_get_dtype_str(self._ptr)
             _check_ptr(dtype_ptr, "get_dtype_str")
-            # ffi.string creates a copy, so we must free the Rust-allocated original.
             self._dtype_str = ffi.string(dtype_ptr).decode('utf-8')
             lib.ztensor_free_string(dtype_ptr)
         return self._dtype_str
@@ -90,9 +117,17 @@ class TensorMetadata:
     @property
     def dtype(self):
         """Returns the numpy dtype for this tensor."""
-        return DTYPE_ZT_TO_NP.get(self.dtype_str)
+        dtype_str = self.dtype_str
+        dt = DTYPE_ZT_TO_NP.get(dtype_str)
+        if dt is None:
+            if dtype_str == 'bfloat16':
+                raise ZTensorError(
+                    "Cannot read 'bfloat16' tensor as NumPy array because the 'ml_dtypes' "
+                    "package is not installed. Please install it to proceed."
+                )
+            raise ZTensorError(f"Unsupported or unknown dtype string '{dtype_str}' found in tensor metadata.")
+        return dt
 
-    # RE-ENABLED: This property now works because the underlying FFI functions are available.
     @property
     def shape(self):
         if self._shape is None:
@@ -101,7 +136,6 @@ class TensorMetadata:
                 shape_data_ptr = lib.ztensor_metadata_get_shape_data(self._ptr)
                 _check_ptr(shape_data_ptr, "get_shape_data")
                 self._shape = tuple(shape_data_ptr[i] for i in range(shape_len))
-                # Free the array that was allocated on the Rust side.
                 lib.ztensor_free_u64_array(shape_data_ptr, shape_len)
             else:
                 self._shape = tuple()
@@ -115,15 +149,12 @@ class Reader:
         path_bytes = file_path.encode('utf-8')
         ptr = lib.ztensor_reader_open(path_bytes)
         _check_ptr(ptr, f"Reader open: {file_path}")
-        # The pointer is automatically garbage collected by CFFI.
         self._ptr = ffi.gc(ptr, lib.ztensor_reader_free)
 
     def __enter__(self):
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        # CFFI's garbage collector handles freeing the reader pointer automatically.
-        # No explicit free is needed here, simplifying the context manager.
         self._ptr = None
 
     def get_metadata(self, name: str) -> TensorMetadata:
@@ -134,8 +165,21 @@ class Reader:
         _check_ptr(meta_ptr, f"get_metadata: {name}")
         return TensorMetadata(meta_ptr)
 
-    def read_tensor(self, name: str) -> np.ndarray:
-        """Reads a tensor by name and returns it as a NumPy array (zero-copy)."""
+    def read_tensor(self, name: str, to: str = 'numpy'):
+        """
+        Reads a tensor by name and returns it as a NumPy array or PyTorch tensor.
+        This is a zero-copy operation for both formats (for CPU tensors).
+
+        Args:
+            name (str): The name of the tensor to read.
+            to (str): The desired output format. Either 'numpy' (default) or 'torch'.
+
+        Returns:
+            np.ndarray or torch.Tensor: The tensor data.
+        """
+        if to not in ['numpy', 'torch']:
+            raise ValueError(f"Unsupported format: '{to}'. Choose 'numpy' or 'torch'.")
+
         metadata = self.get_metadata(name)
         view_ptr = lib.ztensor_reader_read_tensor_view(self._ptr, metadata._ptr)
         _check_ptr(view_ptr, f"read_tensor: {name}")
@@ -143,15 +187,37 @@ class Reader:
         # Let CFFI manage the lifetime of the view pointer.
         view_ptr = ffi.gc(view_ptr, lib.ztensor_free_tensor_view)
 
-        # CORRECTED: Create array using the subclass, which handles reshaping and memory.
-        array = _ZTensorView(
-            buffer=ffi.buffer(view_ptr.data, view_ptr.len),
-            dtype=metadata.dtype,
-            shape=metadata.shape,
-            view_ptr=view_ptr
-        )
+        if to == 'numpy':
+            # Use the custom _ZTensorView to safely manage the FFI pointer lifetime.
+            return _ZTensorView(
+                buffer=ffi.buffer(view_ptr.data, view_ptr.len),
+                dtype=metadata.dtype,  # This property raises on unsupported dtypes
+                shape=metadata.shape,
+                view_ptr=view_ptr
+            )
+
+        elif to == 'torch':
+            if not TORCH_AVAILABLE:
+                raise ZTensorError("PyTorch is not installed. Cannot return a torch tensor.")
+
+            # Get the corresponding torch dtype, raising if not supported.
+            torch_dtype = DTYPE_ZT_TO_TORCH.get(metadata.dtype_str)
+            if torch_dtype is None:
+                raise ZTensorError(
+                    f"Cannot read tensor '{name}' as a PyTorch tensor. "
+                    f"The dtype '{metadata.dtype_str}' is not supported by PyTorch."
+                )
+
+            # Create a tensor directly from the buffer to avoid numpy conversion issues.
+            buffer = ffi.buffer(view_ptr.data, view_ptr.len)
+            torch_tensor = torch.frombuffer(buffer, dtype=torch_dtype).reshape(metadata.shape)
 
-        return array
+            # CRITICAL: Attach the memory owner to the tensor to manage its lifetime.
+            # This ensures the Rust memory (held by view_ptr) is not freed while the
+            # torch tensor is still in use.
+            torch_tensor._owner = view_ptr
+
+            return torch_tensor
 
 
 class Writer:
@@ -161,8 +227,6 @@ class Writer:
         path_bytes = file_path.encode('utf-8')
         ptr = lib.ztensor_writer_create(path_bytes)
         _check_ptr(ptr, f"Writer create: {file_path}")
-        # The pointer is consumed by finalize, so we don't use ffi.gc here.
-        # The writer should be freed via finalize or ztensor_writer_free if finalize fails.
         self._ptr = ptr
         self._finalized = False
 
@@ -170,45 +234,67 @@ class Writer:
         return self
 
     def __exit__(self, exc_type, exc_val, exc_tb):
-        # Automatically finalize on exit if not already done and no error occurred.
         if self._ptr and not self._finalized:
             if exc_type is None:
                 self.finalize()
             else:
-                # If an error occurred, don't finalize, just free the writer to prevent leaks.
                 lib.ztensor_writer_free(self._ptr)
                 self._ptr = None
 
-    def add_tensor(self, name: str, tensor: np.ndarray):
-        """Adds a NumPy array as a tensor to the file."""
-        if not self._ptr: raise ZTensorError("Writer is closed or finalized.")
+    def add_tensor(self, name: str, tensor):
+        """
+        Adds a NumPy or PyTorch tensor to the file (zero-copy).
+        Supports float16 and bfloat16 types.
 
-        name_bytes = name.encode('utf-8')
-        tensor = np.ascontiguousarray(tensor)  # Ensure data is contiguous.
+        Args:
+            name (str): The name of the tensor to add.
+            tensor (np.ndarray or torch.Tensor): The tensor data to write.
+        """
+        if not self._ptr: raise ZTensorError("Writer is closed or finalized.")
 
-        shape_array = np.array(tensor.shape, dtype=np.uint64)
-        shape_ptr = ffi.cast("uint64_t*", shape_array.ctypes.data)
+        # --- Polymorphic tensor handling ---
+        if isinstance(tensor, np.ndarray):
+            tensor = np.ascontiguousarray(tensor)
+            shape = tensor.shape
+            dtype_str = DTYPE_NP_TO_ZT.get(tensor.dtype)
+            data_ptr = ffi.cast("unsigned char*", tensor.ctypes.data)
+            nbytes = tensor.nbytes
+
+        elif TORCH_AVAILABLE and isinstance(tensor, torch.Tensor):
+            if tensor.is_cuda:
+                raise ZTensorError("Cannot write directly from a CUDA tensor. Copy to CPU first using .cpu().")
+            tensor = tensor.contiguous()
+            shape = tuple(tensor.shape)
+            dtype_str = DTYPE_TORCH_TO_ZT.get(tensor.dtype)
+            data_ptr = ffi.cast("unsigned char*", tensor.data_ptr())
+            nbytes = tensor.numel() * tensor.element_size()
+
+        else:
+            supported = "np.ndarray" + (" or torch.Tensor" if TORCH_AVAILABLE else "")
+            raise TypeError(f"Unsupported tensor type: {type(tensor)}. Must be {supported}.")
 
-        dtype_str = DTYPE_NP_TO_ZT.get(tensor.dtype)
         if not dtype_str:
-            raise ZTensorError(f"Unsupported NumPy dtype: {tensor.dtype}")
-        dtype_bytes = dtype_str.encode('utf-8')
+            msg = f"Unsupported dtype: {tensor.dtype}."
+            if 'bfloat16' in str(tensor.dtype) and not ML_DTYPES_AVAILABLE:
+                msg += " For NumPy bfloat16 support, please install the 'ml_dtypes' package."
+            raise ZTensorError(msg)
 
-        # CORRECTED: Cast to `unsigned char*` to match the CFFI definition and Rust FFI.
-        data_ptr = ffi.cast("unsigned char*", tensor.ctypes.data)
+        name_bytes = name.encode('utf-8')
+        shape_array = np.array(shape, dtype=np.uint64)
+        shape_ptr = ffi.cast("uint64_t*", shape_array.ctypes.data)
+        dtype_bytes = dtype_str.encode('utf-8')
 
         status = lib.ztensor_writer_add_tensor(
-            self._ptr, name_bytes, shape_ptr, len(tensor.shape),
-            dtype_bytes, data_ptr, tensor.nbytes
+            self._ptr, name_bytes, shape_ptr, len(shape),
+            dtype_bytes, data_ptr, nbytes
         )
         _check_status(status, f"add_tensor: {name}")
 
     def finalize(self):
         """Finalizes the zTensor file, writing the metadata index."""
         if not self._ptr: raise ZTensorError("Writer is already closed or finalized.")
-
         status = lib.ztensor_writer_finalize(self._ptr)
-        self._ptr = None  # The writer pointer is consumed and invalidated by the Rust call.
+        self._ptr = None
         self._finalized = True
         _check_status(status, "finalize")