ztensor 0.1.0__py3-none-musllinux_1_2_armv7l.whl → 0.1.1__py3-none-musllinux_1_2_armv7l.whl

ztensor/__init__.py

@@ -1,5 +1,6 @@
 import numpy as np
 from .ztensor import ffi, lib
+from ml_dtypes import bfloat16
 
 
 # --- Pythonic Wrapper ---
@@ -52,7 +53,10 @@ DTYPE_NP_TO_ZT = {
     np.dtype('uint64'): 'uint64', np.dtype('uint32'): 'uint32',
     np.dtype('uint16'): 'uint16', np.dtype('uint8'): 'uint8',
     np.dtype('bool'): 'bool',
+    # ADDED: Mapping for bfloat16 to handle writing
+    np.dtype(bfloat16): 'bfloat16',
 }
+# MODIFIED: Re-construct the reverse mapping
 DTYPE_ZT_TO_NP = {v: k for k, v in DTYPE_NP_TO_ZT.items()}
 
 
@@ -90,6 +94,7 @@ class TensorMetadata:
     @property
     def dtype(self):
         """Returns the numpy dtype for this tensor."""
+        # This will now correctly return a bfloat16 dtype object when dtype_str is 'bfloat16'
         return DTYPE_ZT_TO_NP.get(self.dtype_str)
 
     # RE-ENABLED: This property now works because the underlying FFI functions are available.
@@ -137,6 +142,12 @@ class Reader:
     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)
+
+        # This check is now more robust.
+        dtype = metadata.dtype
+        if dtype is None:
+            raise ZTensorError(f"Unsupported ztensor dtype: '{metadata.dtype_str}'")
+
         view_ptr = lib.ztensor_reader_read_tensor_view(self._ptr, metadata._ptr)
         _check_ptr(view_ptr, f"read_tensor: {name}")
 
@@ -144,9 +155,11 @@ class Reader:
         view_ptr = ffi.gc(view_ptr, lib.ztensor_free_tensor_view)
 
         # CORRECTED: Create array using the subclass, which handles reshaping and memory.
+        # With the correct bfloat16 dtype, np.frombuffer will now correctly interpret
+        # the buffer, creating an array with the right number of elements (2048).
         array = _ZTensorView(
             buffer=ffi.buffer(view_ptr.data, view_ptr.len),
-            dtype=metadata.dtype,
+            dtype=dtype,
             shape=metadata.shape,
             view_ptr=view_ptr
         )
@@ -189,6 +202,7 @@ class Writer:
         shape_array = np.array(tensor.shape, dtype=np.uint64)
         shape_ptr = ffi.cast("uint64_t*", shape_array.ctypes.data)
 
+        # The updated DTYPE_NP_TO_ZT will now correctly handle bfloat16 tensors.
         dtype_str = DTYPE_NP_TO_ZT.get(tensor.dtype)
         if not dtype_str:
             raise ZTensorError(f"Unsupported NumPy dtype: {tensor.dtype}")

{ztensor-0.1.0.dist-info → ztensor-0.1.1.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: ztensor
-Version: 0.1.0
+Version: 0.1.1
 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.dist-info → ztensor-0.1.1.dist-info}/RECORD

@@ -1,9 +1,9 @@
-ztensor-0.1.0.dist-info/METADATA,sha256=PE8RgtYKK6sfQK8drIfswyyZxYCHwT1J2Iw1WR6vjtA,4421
-ztensor-0.1.0.dist-info/WHEEL,sha256=uk6lXffjnM20NjdJS1j5pULyv9smQtZKAlwiwmTYIRE,104
-ztensor-0.1.0.dist-info/licenses/LICENSE,sha256=AoeyV1LzTyOz9sbr6uOzk_P0lW963DvhJHnVNVQlI3Y,1063
+ztensor-0.1.1.dist-info/METADATA,sha256=dpaZnxqi7u30Z3uH641XoVg_DeH5k25UamSBjzseBos,4446
+ztensor-0.1.1.dist-info/WHEEL,sha256=uk6lXffjnM20NjdJS1j5pULyv9smQtZKAlwiwmTYIRE,104
+ztensor-0.1.1.dist-info/licenses/LICENSE,sha256=AoeyV1LzTyOz9sbr6uOzk_P0lW963DvhJHnVNVQlI3Y,1063
 ztensor.libs/libgcc_s-5b5488a6.so.1,sha256=HGKUsVmTeNAxEdSy7Ua5Vh_I9FN3RCbPWzvZ7H_TrwE,2749061
-ztensor/__init__.py,sha256=sr2eKuzJBZLDbqIAntJ6WDxSZLBfWoi9_UN464jf8lw,8403
+ztensor/__init__.py,sha256=pzpS0XNcpBb762OPUG7-3raEJxz7P99KX6CGAdh-DS0,9086
 ztensor/ztensor/__init__.py,sha256=sIpB0pJYFX20TdZapIoPMxqMz37wKJyxCkAlTemWDq4,140
 ztensor/ztensor/ffi.py,sha256=J7CG26lx0Xu0IjYR7GWev-BDgGQXseKA4cjkh0IhnLE,2746
 ztensor/ztensor/libztensor.so,sha256=9bgmcP7Tj8o3loKKYgTfK6kzL1wcgxoKCTykU6CSvEI,1611249
-ztensor-0.1.0.dist-info/RECORD,,
+ztensor-0.1.1.dist-info/RECORD,,