ztensor 0.1.0__py3-none-manylinux_2_5_i686.manylinux1_i686.whl

ztensor/__init__.py

@@ -0,0 +1,216 @@
+import numpy as np
+from .ztensor import ffi, lib
+
+
+# --- Pythonic Wrapper ---
+class ZTensorError(Exception):
+    """Custom exception for ztensor-related errors."""
+    pass
+
+
+# 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
+        return obj
+
+    def __array_finalize__(self, obj):
+        # This ensures that views and slices of our array also hold the reference.
+        if obj is None: return
+        self._owner = getattr(obj, '_owner', None)
+
+
+def _get_last_error():
+    """Retrieves the last error message from the Rust library."""
+    err_msg_ptr = lib.ztensor_last_error_message()
+    if err_msg_ptr != ffi.NULL:
+        return ffi.string(err_msg_ptr).decode('utf-8')
+    return "Unknown FFI error"
+
+
+def _check_ptr(ptr, func_name=""):
+    """Checks if a pointer from the FFI is null and raises an error if it is."""
+    if ptr == ffi.NULL:
+        raise ZTensorError(f"Error in {func_name}: {_get_last_error()}")
+    return ptr
+
+
+def _check_status(status, func_name=""):
+    """Checks the integer status code from an FFI call and raises on failure."""
+    if status != 0:
+        raise ZTensorError(f"Error in {func_name}: {_get_last_error()}")
+
+
+# Type Mappings between NumPy and ztensor
+DTYPE_NP_TO_ZT = {
+    np.dtype('float64'): 'float64', np.dtype('float32'): 'float32',
+    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',
+}
+DTYPE_ZT_TO_NP = {v: k for k, v in DTYPE_NP_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
+        self._dtype_str = None
+        self._shape = None
+
+    @property
+    def name(self):
+        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
+
+    @property
+    def dtype_str(self):
+        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
+
+    @property
+    def dtype(self):
+        """Returns the numpy dtype for this tensor."""
+        return DTYPE_ZT_TO_NP.get(self.dtype_str)
+
+    # RE-ENABLED: This property now works because the underlying FFI functions are available.
+    @property
+    def shape(self):
+        if self._shape is None:
+            shape_len = lib.ztensor_metadata_get_shape_len(self._ptr)
+            if shape_len > 0:
+                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()
+        return self._shape
+
+
+class Reader:
+    """A Pythonic context manager for reading zTensor files."""
+
+    def __init__(self, file_path):
+        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:
+        """Retrieves metadata for a tensor by its name."""
+        if self._ptr is None: raise ZTensorError("Reader is closed.")
+        name_bytes = name.encode('utf-8')
+        meta_ptr = lib.ztensor_reader_get_metadata_by_name(self._ptr, name_bytes)
+        _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)."""
+        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}")
+
+        # 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
+        )
+
+        return array
+
+
+class Writer:
+    """A Pythonic context manager for writing zTensor files."""
+
+    def __init__(self, file_path):
+        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
+
+    def __enter__(self):
+        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.")
+
+        name_bytes = name.encode('utf-8')
+        tensor = np.ascontiguousarray(tensor)  # Ensure data is contiguous.
+
+        shape_array = np.array(tensor.shape, dtype=np.uint64)
+        shape_ptr = ffi.cast("uint64_t*", shape_array.ctypes.data)
+
+        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')
+
+        # CORRECTED: Cast to `unsigned char*` to match the CFFI definition and Rust FFI.
+        data_ptr = ffi.cast("unsigned char*", tensor.ctypes.data)
+
+        status = lib.ztensor_writer_add_tensor(
+            self._ptr, name_bytes, shape_ptr, len(tensor.shape),
+            dtype_bytes, data_ptr, tensor.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._finalized = True
+        _check_status(status, "finalize")
+
+
+__all__ = ["Reader", "Writer", "TensorMetadata", "ZTensorError"]

ztensor/ztensor/__init__.py

@@ -0,0 +1,7 @@
+__all__ = ["lib", "ffi"]
+
+import os
+from .ffi import ffi
+
+lib = ffi.dlopen(os.path.join(os.path.dirname(__file__), 'libztensor.so'))
+del os

ztensor/ztensor/ffi.py

@@ -0,0 +1,10 @@
+# auto-generated file
+import _cffi_backend
+
+ffi = _cffi_backend.FFI('ffi',
+    _version = 0x2601,
+    _types = b'\x00\x00\x29\x0D\x00\x00\x3D\x03\x00\x00\x3C\x03\x00\x00\x00\x0F\x00\x00\x2C\x0D\x00\x00\x3D\x03\x00\x00\x3F\x03\x00\x00\x00\x0F\x00\x00\x01\x0D\x00\x00\x06\x11\x00\x00\x00\x0F\x00\x00\x14\x0D\x00\x00\x06\x11\x00\x00\x00\x0F\x00\x00\x35\x0D\x00\x00\x02\x11\x00\x00\x00\x0F\x00\x00\x06\x0D\x00\x00\x00\x0F\x00\x00\x40\x0D\x00\x00\x3E\x03\x00\x00\x00\x0F\x00\x00\x40\x0D\x00\x00\x14\x11\x00\x00\x06\x11\x00\x00\x41\x03\x00\x00\x1C\x01\x00\x00\x06\x11\x00\x00\x42\x03\x00\x00\x1C\x01\x00\x00\x00\x0F\x00\x00\x1A\x0D\x00\x00\x02\x11\x00\x00\x00\x0F\x00\x00\x1A\x0D\x00\x00\x05\x11\x00\x00\x00\x0F\x00\x00\x38\x0D\x00\x00\x02\x11\x00\x00\x00\x0F\x00\x00\x44\x0D\x00\x00\x3B\x03\x00\x00\x00\x0F\x00\x00\x44\x0D\x00\x00\x3C\x03\x00\x00\x00\x0F\x00\x00\x44\x0D\x00\x00\x01\x11\x00\x00\x00\x0F\x00\x00\x44\x0D\x00\x00\x14\x11\x00\x00\x00\x0F\x00\x00\x44\x0D\x00\x00\x3F\x03\x00\x00\x00\x0F\x00\x00\x44\x0D\x00\x00\x41\x03\x00\x00\x1C\x01\x00\x00\x00\x0F\x00\x00\x00\x09\x00\x00\x01\x09\x00\x00\x02\x09\x00\x00\x03\x09\x00\x00\x02\x01\x00\x00\x07\x01\x00\x00\x18\x01\x00\x00\x04\x01\x00\x00\x44\x03\x00\x00\x00\x01',
+    _globals = (b'\xFF\xFF\xFF\x1FALIGNMENT',64,b'\x00\x00\x34\x23ztensor_free_string',0,b'\x00\x00\x28\x23ztensor_free_tensor_view',0,b'\x00\x00\x37\x23ztensor_free_u64_array',0,b'\x00\x00\x11\x23ztensor_last_error_message',0,b'\x00\x00\x2B\x23ztensor_metadata_free',0,b'\x00\x00\x0E\x23ztensor_metadata_get_dtype_str',0,b'\x00\x00\x0E\x23ztensor_metadata_get_name',0,b'\x00\x00\x25\x23ztensor_metadata_get_shape_data',0,b'\x00\x00\x1F\x23ztensor_metadata_get_shape_len',0,b'\x00\x00\x2E\x23ztensor_reader_free',0,b'\x00\x00\x04\x23ztensor_reader_get_metadata_by_name',0,b'\x00\x00\x22\x23ztensor_reader_get_metadata_count',0,b'\x00\x00\x08\x23ztensor_reader_open',0,b'\x00\x00\x00\x23ztensor_reader_read_tensor_view',0,b'\x00\x00\x16\x23ztensor_writer_add_tensor',0,b'\x00\x00\x0B\x23ztensor_writer_create',0,b'\x00\x00\x13\x23ztensor_writer_finalize',0,b'\x00\x00\x31\x23ztensor_writer_free',0),
+    _struct_unions = ((b'\x00\x00\x00\x3B\x00\x00\x00\x02CTensorDataView',b'\x00\x00\x1C\x11data',b'\x00\x00\x1A\x11len',b'\x00\x00\x43\x11_owner'),(b'\x00\x00\x00\x3C\x00\x00\x00\x10TensorMetadata',),(b'\x00\x00\x00\x3D\x00\x00\x00\x10ZTensorReader_BufReader_File',),(b'\x00\x00\x00\x3E\x00\x00\x00\x10ZTensorWriter_BufWriter_File',)),
+    _typenames = (b'\x00\x00\x00\x3BCTensorDataView',b'\x00\x00\x00\x3CCTensorMetadata',b'\x00\x00\x00\x3DCZTensorReader',b'\x00\x00\x00\x3ECZTensorWriter',b'\x00\x00\x00\x3CTensorMetadata',b'\x00\x00\x00\x3DZTensorReader_BufReader_File',b'\x00\x00\x00\x3EZTensorWriter_BufWriter_File'),
+)

ztensor-0.1.0.dist-info/METADATA

@@ -0,0 +1,105 @@
+Metadata-Version: 2.4
+Name: ztensor
+Version: 0.1.0
+Classifier: Programming Language :: Rust
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Operating System :: OS Independent
+Classifier: Intended Audience :: Developers
+Classifier: Topic :: Scientific/Engineering
+Requires-Dist: numpy
+Requires-Dist: cffi
+License-File: LICENSE
+Summary: Python bindings for the zTensor library.
+Author: In Gim <in.gim@yale.edu>
+Author-email: In Gim <in.gim@yale.edu>
+License: MIT
+Description-Content-Type: text/markdown; charset=UTF-8; variant=GFM
+Project-URL: Homepage, https://github.com/pie-project/ztensor
+Project-URL: Source, https://github.com/pie-project/ztensor
+
+# zTensor File Format
+
+**Version 0.1.0**
+
+zTensor is a binary format for storing large multi-dimensional arrays (tensors), designed for efficient, safe, and flexible access. It supports raw and compressed (zstd) encodings, quantized and sparse layouts, and is extensible.
+
+For dense tensors stored with encoding: "raw" and matching endianness, zTensor enables zero-copy access to tensor data.
+
+## File Layout
+
+```
++-------------------------------+
+| Magic Number (8 bytes)        |
++-------------------------------+
+| Tensor Blob 0 (aligned)       |
++-------------------------------+
+| Padding (if needed)           |
++-------------------------------+
+| Tensor Blob 1 (aligned)       |
++-------------------------------+
+| ...                           |
++-------------------------------+
+| CBOR Metadata Array           |
++-------------------------------+
+| CBOR Array Size (8 bytes)     |
++-------------------------------+
+```
+
+- **Magic Number:** ASCII "ZTEN0001" at offset 0.
+- **Tensor Blobs:** Each tensor's data, starting at a 64-byte aligned offset. No per-blob headers. Padding (undefined value, usually zero) is inserted as needed.
+- **CBOR Metadata Array:** At the end of the file, a CBOR-encoded array of metadata maps (one per tensor).
+- **CBOR Array Size:** Last 8 bytes, little-endian uint64, gives the size of the CBOR array.
+
+## Tensor Metadata (CBOR Map)
+Each tensor's metadata is a CBOR map with these required fields:
+- `name` (string): Tensor name.
+- `offset` (uint64): Absolute file offset to tensor data (multiple of 64).
+- `size` (uint64): On-disk size in bytes (compressed if applicable).
+- `dtype` (string): Data type (see below).
+- `shape` (array): Array of dimensions (empty for scalar).
+- `encoding` (string): Data encoding (see below).
+- `layout` (string): "dense" (default) or "sparse". For sparse, see below.
+
+Optional fields:
+- `data_endianness` (string): "little" or "big" (default: little, for raw multi-byte types).
+- `checksum` (string): e.g., "crc32c:0x1234ABCD" or "sha256:...".
+- Custom fields are allowed; unknown keys are ignored by readers.
+
+## Supported Data Types (`dtype`)
+- `float64`, `float32`, `float16`, `bfloat16`
+- `int64`, `int32`, `int16`, `int8`
+- `uint64`, `uint32`, `uint16`, `uint8`
+- `bool`
+
+## Supported Encodings (`encoding`)
+- `raw`: Direct binary dump of tensor elements (with `data_endianness` if multi-byte).
+- `zstd`: Zstandard-compressed data. `size` is compressed size.
+
+## Layouts
+- `dense` (default): Standard contiguous tensor data.
+- `sparse`: Data is stored in a sparse format. The metadata map must include a `sparse_format` field (e.g., "csr", "coo") and any additional fields required to describe the sparse structure (such as index arrays, indptr, etc.).
+
+## Index Reading
+To read the index:
+1. Read the last 8 bytes for the CBOR array size.
+2. Seek backwards by that amount to read the CBOR metadata array.
+
+As a result, the start offset of the metadata is: (file size) - (size of the metadata) - (8 byte).
+
+## Zero-Tensor Files
+A valid zTensor file may contain zero tensors:
+- 8 bytes magic, 1 byte empty CBOR array (`0x80`), 8 bytes size (`0x01...00`).
+- Total: 17 bytes.
+
+## Extensibility
+- New `dtype`, `encoding`, and `layout` values may be added in future versions.
+- Custom metadata fields are allowed; unknown fields are ignored.
+
+## Notes
+- All offsets are absolute and account for the 8-byte magic number.
+- All tensor data blobs are 64-byte aligned.
+- No per-blob headers; all metadata is in the CBOR array at the end of the file.
+- For `encoding: "raw"` and multi-byte `dtype`, data is little-endian unless `data_endianness` is specified.
+- For sparse tensors, the metadata must fully describe the sparse structure.
+

ztensor-0.1.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+ztensor-0.1.0.dist-info/METADATA,sha256=PE8RgtYKK6sfQK8drIfswyyZxYCHwT1J2Iw1WR6vjtA,4421
+ztensor-0.1.0.dist-info/WHEEL,sha256=SpGB6IWzB_riHhA8KE-DtiG1EenXfRr8HYaf2B05sKc,118
+ztensor-0.1.0.dist-info/licenses/LICENSE,sha256=AoeyV1LzTyOz9sbr6uOzk_P0lW963DvhJHnVNVQlI3Y,1063
+ztensor/__init__.py,sha256=sr2eKuzJBZLDbqIAntJ6WDxSZLBfWoi9_UN464jf8lw,8403
+ztensor/ztensor/__init__.py,sha256=sIpB0pJYFX20TdZapIoPMxqMz37wKJyxCkAlTemWDq4,140
+ztensor/ztensor/ffi.py,sha256=J7CG26lx0Xu0IjYR7GWev-BDgGQXseKA4cjkh0IhnLE,2746
+ztensor/ztensor/libztensor.so,sha256=UQ1w5wToRUBAf80aoZtyzgS6DuM0WM7PfmlJ2qi238I,1799904
+ztensor-0.1.0.dist-info/RECORD,,

ztensor-0.1.0.dist-info/WHEEL

@@ -0,0 +1,4 @@
+Wheel-Version: 1.0
+Generator: maturin (1.9.0)
+Root-Is-Purelib: false
+Tag: py3-none-manylinux_2_5_i686.manylinux1_i686

ztensor-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 In Gim
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.