ztensor 0.1.1__py3-none-musllinux_1_2_armv7l.whl → 0.1.4__py3-none-musllinux_1_2_armv7l.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Potentially problematic release.
This version of ztensor might be problematic. Click here for more details.
- ztensor/__init__.py +234 -56
- ztensor/ztensor/ffi.py +4 -4
- ztensor/ztensor/libztensor.so +0 -0
- {ztensor-0.1.1.dist-info → ztensor-0.1.4.dist-info}/METADATA +1 -1
- ztensor-0.1.4.dist-info/RECORD +9 -0
- {ztensor-0.1.1.dist-info → ztensor-0.1.4.dist-info}/WHEEL +1 -1
- ztensor-0.1.1.dist-info/RECORD +0 -9
- {ztensor-0.1.1.dist-info → ztensor-0.1.4.dist-info}/licenses/LICENSE +0 -0
ztensor/__init__.py
CHANGED
|
@@ -1,6 +1,22 @@
|
|
|
1
1
|
import numpy as np
|
|
2
2
|
from .ztensor import ffi, lib
|
|
3
|
-
|
|
3
|
+
|
|
4
|
+
# --- Optional PyTorch Import ---
|
|
5
|
+
try:
|
|
6
|
+
import torch
|
|
7
|
+
|
|
8
|
+
TORCH_AVAILABLE = True
|
|
9
|
+
except ImportError:
|
|
10
|
+
TORCH_AVAILABLE = False
|
|
11
|
+
|
|
12
|
+
# --- Optional ml_dtypes for bfloat16 in NumPy ---
|
|
13
|
+
try:
|
|
14
|
+
from ml_dtypes import bfloat16 as np_bfloat16
|
|
15
|
+
|
|
16
|
+
ML_DTYPES_AVAILABLE = True
|
|
17
|
+
except ImportError:
|
|
18
|
+
np_bfloat16 = None
|
|
19
|
+
ML_DTYPES_AVAILABLE = False
|
|
4
20
|
|
|
5
21
|
|
|
6
22
|
# --- Pythonic Wrapper ---
|
|
@@ -12,7 +28,6 @@ class ZTensorError(Exception):
|
|
|
12
28
|
# A custom ndarray subclass to safely manage the lifetime of the CFFI pointer.
|
|
13
29
|
class _ZTensorView(np.ndarray):
|
|
14
30
|
def __new__(cls, buffer, dtype, shape, view_ptr):
|
|
15
|
-
# Create an array from the buffer, reshape it, and cast it to our custom type.
|
|
16
31
|
obj = np.frombuffer(buffer, dtype=dtype).reshape(shape).view(cls)
|
|
17
32
|
# Attach the object that owns the memory to an attribute.
|
|
18
33
|
obj._owner = view_ptr
|
|
@@ -45,73 +60,158 @@ def _check_status(status, func_name=""):
|
|
|
45
60
|
raise ZTensorError(f"Error in {func_name}: {_get_last_error()}")
|
|
46
61
|
|
|
47
62
|
|
|
48
|
-
# Type Mappings
|
|
63
|
+
# --- Type Mappings ---
|
|
64
|
+
# NumPy Mappings
|
|
49
65
|
DTYPE_NP_TO_ZT = {
|
|
50
|
-
np.dtype('float64'): 'float64', np.dtype('float32'): 'float32',
|
|
66
|
+
np.dtype('float64'): 'float64', np.dtype('float32'): 'float32', np.dtype('float16'): 'float16',
|
|
51
67
|
np.dtype('int64'): 'int64', np.dtype('int32'): 'int32',
|
|
52
68
|
np.dtype('int16'): 'int16', np.dtype('int8'): 'int8',
|
|
53
69
|
np.dtype('uint64'): 'uint64', np.dtype('uint32'): 'uint32',
|
|
54
70
|
np.dtype('uint16'): 'uint16', np.dtype('uint8'): 'uint8',
|
|
55
71
|
np.dtype('bool'): 'bool',
|
|
56
|
-
# ADDED: Mapping for bfloat16 to handle writing
|
|
57
|
-
np.dtype(bfloat16): 'bfloat16',
|
|
58
72
|
}
|
|
59
|
-
|
|
73
|
+
if ML_DTYPES_AVAILABLE:
|
|
74
|
+
DTYPE_NP_TO_ZT[np.dtype(np_bfloat16)] = 'bfloat16'
|
|
60
75
|
DTYPE_ZT_TO_NP = {v: k for k, v in DTYPE_NP_TO_ZT.items()}
|
|
61
76
|
|
|
77
|
+
# PyTorch Mappings (if available)
|
|
78
|
+
if TORCH_AVAILABLE:
|
|
79
|
+
DTYPE_TORCH_TO_ZT = {
|
|
80
|
+
torch.float64: 'float64', torch.float32: 'float32', torch.float16: 'float16',
|
|
81
|
+
torch.bfloat16: 'bfloat16',
|
|
82
|
+
torch.int64: 'int64', torch.int32: 'int32',
|
|
83
|
+
torch.int16: 'int16', torch.int8: 'int8',
|
|
84
|
+
torch.uint8: 'uint8', torch.bool: 'bool',
|
|
85
|
+
}
|
|
86
|
+
DTYPE_ZT_TO_TORCH = {v: k for k, v in DTYPE_TORCH_TO_ZT.items()}
|
|
87
|
+
|
|
62
88
|
|
|
63
89
|
class TensorMetadata:
|
|
64
90
|
"""A Pythonic wrapper around the CTensorMetadata pointer."""
|
|
65
91
|
|
|
66
92
|
def __init__(self, meta_ptr):
|
|
67
|
-
# The pointer is now automatically garbage collected by CFFI when this object dies.
|
|
68
93
|
self._ptr = ffi.gc(meta_ptr, lib.ztensor_metadata_free)
|
|
69
94
|
_check_ptr(self._ptr, "TensorMetadata constructor")
|
|
95
|
+
# Cache for properties to avoid repeated FFI calls
|
|
70
96
|
self._name = None
|
|
71
97
|
self._dtype_str = None
|
|
72
98
|
self._shape = None
|
|
99
|
+
self._offset = None
|
|
100
|
+
self._size = None
|
|
101
|
+
self._layout = None
|
|
102
|
+
self._encoding = None
|
|
103
|
+
self._endianness = "not_checked"
|
|
104
|
+
self._checksum = "not_checked"
|
|
105
|
+
|
|
106
|
+
def __repr__(self):
|
|
107
|
+
return f"<TensorMetadata name='{self.name}' shape={self.shape} dtype='{self.dtype_str}'>"
|
|
73
108
|
|
|
74
109
|
@property
|
|
75
110
|
def name(self):
|
|
111
|
+
"""The name of the tensor."""
|
|
76
112
|
if self._name is None:
|
|
77
113
|
name_ptr = lib.ztensor_metadata_get_name(self._ptr)
|
|
78
114
|
_check_ptr(name_ptr, "get_name")
|
|
79
|
-
# ffi.string creates a copy, so we must free the Rust-allocated original.
|
|
80
115
|
self._name = ffi.string(name_ptr).decode('utf-8')
|
|
81
116
|
lib.ztensor_free_string(name_ptr)
|
|
82
117
|
return self._name
|
|
83
118
|
|
|
84
119
|
@property
|
|
85
120
|
def dtype_str(self):
|
|
121
|
+
"""The zTensor dtype string (e.g., 'float32')."""
|
|
86
122
|
if self._dtype_str is None:
|
|
87
123
|
dtype_ptr = lib.ztensor_metadata_get_dtype_str(self._ptr)
|
|
88
124
|
_check_ptr(dtype_ptr, "get_dtype_str")
|
|
89
|
-
# ffi.string creates a copy, so we must free the Rust-allocated original.
|
|
90
125
|
self._dtype_str = ffi.string(dtype_ptr).decode('utf-8')
|
|
91
126
|
lib.ztensor_free_string(dtype_ptr)
|
|
92
127
|
return self._dtype_str
|
|
93
128
|
|
|
94
129
|
@property
|
|
95
130
|
def dtype(self):
|
|
96
|
-
"""
|
|
97
|
-
|
|
98
|
-
|
|
131
|
+
"""The numpy dtype for this tensor."""
|
|
132
|
+
dtype_str = self.dtype_str
|
|
133
|
+
dt = DTYPE_ZT_TO_NP.get(dtype_str)
|
|
134
|
+
if dt is None:
|
|
135
|
+
if dtype_str == 'bfloat16':
|
|
136
|
+
raise ZTensorError(
|
|
137
|
+
"Cannot read 'bfloat16' tensor as NumPy array because the 'ml_dtypes' "
|
|
138
|
+
"package is not installed. Please install it to proceed."
|
|
139
|
+
)
|
|
140
|
+
raise ZTensorError(f"Unsupported or unknown dtype string '{dtype_str}' found in tensor metadata.")
|
|
141
|
+
return dt
|
|
99
142
|
|
|
100
|
-
# RE-ENABLED: This property now works because the underlying FFI functions are available.
|
|
101
143
|
@property
|
|
102
144
|
def shape(self):
|
|
145
|
+
"""The shape of the tensor as a tuple."""
|
|
103
146
|
if self._shape is None:
|
|
104
147
|
shape_len = lib.ztensor_metadata_get_shape_len(self._ptr)
|
|
105
148
|
if shape_len > 0:
|
|
106
149
|
shape_data_ptr = lib.ztensor_metadata_get_shape_data(self._ptr)
|
|
107
150
|
_check_ptr(shape_data_ptr, "get_shape_data")
|
|
108
151
|
self._shape = tuple(shape_data_ptr[i] for i in range(shape_len))
|
|
109
|
-
# Free the array that was allocated on the Rust side.
|
|
110
152
|
lib.ztensor_free_u64_array(shape_data_ptr, shape_len)
|
|
111
153
|
else:
|
|
112
154
|
self._shape = tuple()
|
|
113
155
|
return self._shape
|
|
114
156
|
|
|
157
|
+
@property
|
|
158
|
+
def offset(self):
|
|
159
|
+
"""The on-disk offset of the tensor data in bytes."""
|
|
160
|
+
if self._offset is None:
|
|
161
|
+
self._offset = lib.ztensor_metadata_get_offset(self._ptr)
|
|
162
|
+
return self._offset
|
|
163
|
+
|
|
164
|
+
@property
|
|
165
|
+
def size(self):
|
|
166
|
+
"""The on-disk size of the tensor data in bytes (can be compressed size)."""
|
|
167
|
+
if self._size is None:
|
|
168
|
+
self._size = lib.ztensor_metadata_get_size(self._ptr)
|
|
169
|
+
return self._size
|
|
170
|
+
|
|
171
|
+
@property
|
|
172
|
+
def layout(self):
|
|
173
|
+
"""The tensor layout as a string (e.g., 'dense')."""
|
|
174
|
+
if self._layout is None:
|
|
175
|
+
layout_ptr = lib.ztensor_metadata_get_layout_str(self._ptr)
|
|
176
|
+
_check_ptr(layout_ptr, "get_layout_str")
|
|
177
|
+
self._layout = ffi.string(layout_ptr).decode('utf-8')
|
|
178
|
+
lib.ztensor_free_string(layout_ptr)
|
|
179
|
+
return self._layout
|
|
180
|
+
|
|
181
|
+
@property
|
|
182
|
+
def encoding(self):
|
|
183
|
+
"""The tensor encoding as a string (e.g., 'raw', 'zstd')."""
|
|
184
|
+
if self._encoding is None:
|
|
185
|
+
encoding_ptr = lib.ztensor_metadata_get_encoding_str(self._ptr)
|
|
186
|
+
_check_ptr(encoding_ptr, "get_encoding_str")
|
|
187
|
+
self._encoding = ffi.string(encoding_ptr).decode('utf-8')
|
|
188
|
+
lib.ztensor_free_string(encoding_ptr)
|
|
189
|
+
return self._encoding
|
|
190
|
+
|
|
191
|
+
@property
|
|
192
|
+
def endianness(self):
|
|
193
|
+
"""The data endianness ('little', 'big') if applicable, else None."""
|
|
194
|
+
if self._endianness == "not_checked":
|
|
195
|
+
endian_ptr = lib.ztensor_metadata_get_data_endianness_str(self._ptr)
|
|
196
|
+
if endian_ptr == ffi.NULL:
|
|
197
|
+
self._endianness = None
|
|
198
|
+
else:
|
|
199
|
+
self._endianness = ffi.string(endian_ptr).decode('utf-8')
|
|
200
|
+
lib.ztensor_free_string(endian_ptr)
|
|
201
|
+
return self._endianness
|
|
202
|
+
|
|
203
|
+
@property
|
|
204
|
+
def checksum(self):
|
|
205
|
+
"""The checksum string if present, else None."""
|
|
206
|
+
if self._checksum == "not_checked":
|
|
207
|
+
checksum_ptr = lib.ztensor_metadata_get_checksum_str(self._ptr)
|
|
208
|
+
if checksum_ptr == ffi.NULL:
|
|
209
|
+
self._checksum = None
|
|
210
|
+
else:
|
|
211
|
+
self._checksum = ffi.string(checksum_ptr).decode('utf-8')
|
|
212
|
+
lib.ztensor_free_string(checksum_ptr)
|
|
213
|
+
return self._checksum
|
|
214
|
+
|
|
115
215
|
|
|
116
216
|
class Reader:
|
|
117
217
|
"""A Pythonic context manager for reading zTensor files."""
|
|
@@ -120,17 +220,47 @@ class Reader:
|
|
|
120
220
|
path_bytes = file_path.encode('utf-8')
|
|
121
221
|
ptr = lib.ztensor_reader_open(path_bytes)
|
|
122
222
|
_check_ptr(ptr, f"Reader open: {file_path}")
|
|
123
|
-
# The pointer is automatically garbage collected by CFFI.
|
|
124
223
|
self._ptr = ffi.gc(ptr, lib.ztensor_reader_free)
|
|
125
224
|
|
|
126
225
|
def __enter__(self):
|
|
127
226
|
return self
|
|
128
227
|
|
|
129
228
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
130
|
-
# CFFI's garbage collector handles freeing the reader pointer automatically.
|
|
131
|
-
# No explicit free is needed here, simplifying the context manager.
|
|
132
229
|
self._ptr = None
|
|
133
230
|
|
|
231
|
+
def __len__(self):
|
|
232
|
+
"""Returns the number of tensors in the file."""
|
|
233
|
+
if self._ptr is None: raise ZTensorError("Reader is closed.")
|
|
234
|
+
return lib.ztensor_reader_get_metadata_count(self._ptr)
|
|
235
|
+
|
|
236
|
+
def __iter__(self):
|
|
237
|
+
"""Iterates over the metadata of all tensors in the file."""
|
|
238
|
+
if self._ptr is None: raise ZTensorError("Reader is closed.")
|
|
239
|
+
for i in range(len(self)):
|
|
240
|
+
yield self[i]
|
|
241
|
+
|
|
242
|
+
def __getitem__(self, index: int) -> TensorMetadata:
|
|
243
|
+
"""Retrieves metadata for a tensor by its index."""
|
|
244
|
+
if self._ptr is None: raise ZTensorError("Reader is closed.")
|
|
245
|
+
if index >= len(self):
|
|
246
|
+
raise IndexError("Tensor index out of range")
|
|
247
|
+
meta_ptr = lib.ztensor_reader_get_metadata_by_index(self._ptr, index)
|
|
248
|
+
_check_ptr(meta_ptr, f"get_metadata_by_index: {index}")
|
|
249
|
+
return TensorMetadata(meta_ptr)
|
|
250
|
+
|
|
251
|
+
def list_tensors(self) -> list[TensorMetadata]:
|
|
252
|
+
"""Returns a list of all TensorMetadata objects in the file."""
|
|
253
|
+
return list(self)
|
|
254
|
+
|
|
255
|
+
def get_tensor_names(self) -> list[str]:
|
|
256
|
+
"""Returns a list of all tensor names in the file."""
|
|
257
|
+
if self._ptr is None: raise ZTensorError("Reader is closed.")
|
|
258
|
+
c_array_ptr = lib.ztensor_reader_get_all_tensor_names(self._ptr)
|
|
259
|
+
_check_ptr(c_array_ptr, "get_all_tensor_names")
|
|
260
|
+
c_array_ptr = ffi.gc(c_array_ptr, lib.ztensor_free_string_array)
|
|
261
|
+
|
|
262
|
+
return [ffi.string(c_array_ptr.strings[i]).decode('utf-8') for i in range(c_array_ptr.len)]
|
|
263
|
+
|
|
134
264
|
def get_metadata(self, name: str) -> TensorMetadata:
|
|
135
265
|
"""Retrieves metadata for a tensor by its name."""
|
|
136
266
|
if self._ptr is None: raise ZTensorError("Reader is closed.")
|
|
@@ -139,32 +269,60 @@ class Reader:
|
|
|
139
269
|
_check_ptr(meta_ptr, f"get_metadata: {name}")
|
|
140
270
|
return TensorMetadata(meta_ptr)
|
|
141
271
|
|
|
142
|
-
def read_tensor(self, name: str
|
|
143
|
-
"""
|
|
144
|
-
|
|
272
|
+
def read_tensor(self, name: str, to: str = 'numpy'):
|
|
273
|
+
"""
|
|
274
|
+
Reads a tensor by name and returns it as a NumPy array or PyTorch tensor.
|
|
275
|
+
This is a zero-copy operation for both formats (for CPU tensors).
|
|
276
|
+
|
|
277
|
+
Args:
|
|
278
|
+
name (str): The name of the tensor to read.
|
|
279
|
+
to (str): The desired output format. Either 'numpy' (default) or 'torch'.
|
|
145
280
|
|
|
146
|
-
|
|
147
|
-
|
|
148
|
-
|
|
149
|
-
|
|
281
|
+
Returns:
|
|
282
|
+
np.ndarray or torch.Tensor: The tensor data.
|
|
283
|
+
"""
|
|
284
|
+
if self._ptr is None: raise ZTensorError("Reader is closed.")
|
|
285
|
+
if to not in ['numpy', 'torch']:
|
|
286
|
+
raise ValueError(f"Unsupported format: '{to}'. Choose 'numpy' or 'torch'.")
|
|
150
287
|
|
|
288
|
+
metadata = self.get_metadata(name)
|
|
151
289
|
view_ptr = lib.ztensor_reader_read_tensor_view(self._ptr, metadata._ptr)
|
|
152
290
|
_check_ptr(view_ptr, f"read_tensor: {name}")
|
|
153
291
|
|
|
154
292
|
# Let CFFI manage the lifetime of the view pointer.
|
|
155
293
|
view_ptr = ffi.gc(view_ptr, lib.ztensor_free_tensor_view)
|
|
156
294
|
|
|
157
|
-
|
|
158
|
-
|
|
159
|
-
|
|
160
|
-
|
|
161
|
-
|
|
162
|
-
|
|
163
|
-
|
|
164
|
-
|
|
165
|
-
|
|
295
|
+
if to == 'numpy':
|
|
296
|
+
# Use the custom _ZTensorView to safely manage the FFI pointer lifetime.
|
|
297
|
+
return _ZTensorView(
|
|
298
|
+
buffer=ffi.buffer(view_ptr.data, view_ptr.len),
|
|
299
|
+
dtype=metadata.dtype, # This property raises on unsupported dtypes
|
|
300
|
+
shape=metadata.shape,
|
|
301
|
+
view_ptr=view_ptr
|
|
302
|
+
)
|
|
303
|
+
|
|
304
|
+
elif to == 'torch':
|
|
305
|
+
if not TORCH_AVAILABLE:
|
|
306
|
+
raise ZTensorError("PyTorch is not installed. Cannot return a torch tensor.")
|
|
166
307
|
|
|
167
|
-
|
|
308
|
+
# Get the corresponding torch dtype, raising if not supported.
|
|
309
|
+
torch_dtype = DTYPE_ZT_TO_TORCH.get(metadata.dtype_str)
|
|
310
|
+
if torch_dtype is None:
|
|
311
|
+
raise ZTensorError(
|
|
312
|
+
f"Cannot read tensor '{name}' as a PyTorch tensor. "
|
|
313
|
+
f"The dtype '{metadata.dtype_str}' is not supported by PyTorch."
|
|
314
|
+
)
|
|
315
|
+
|
|
316
|
+
# Create a tensor directly from the buffer to avoid numpy conversion issues.
|
|
317
|
+
buffer = ffi.buffer(view_ptr.data, view_ptr.len)
|
|
318
|
+
torch_tensor = torch.frombuffer(buffer, dtype=torch_dtype).reshape(metadata.shape)
|
|
319
|
+
|
|
320
|
+
# CRITICAL: Attach the memory owner to the tensor to manage its lifetime.
|
|
321
|
+
# This ensures the Rust memory (held by view_ptr) is not freed while the
|
|
322
|
+
# torch tensor is still in use.
|
|
323
|
+
torch_tensor._owner = view_ptr
|
|
324
|
+
|
|
325
|
+
return torch_tensor
|
|
168
326
|
|
|
169
327
|
|
|
170
328
|
class Writer:
|
|
@@ -174,8 +332,6 @@ class Writer:
|
|
|
174
332
|
path_bytes = file_path.encode('utf-8')
|
|
175
333
|
ptr = lib.ztensor_writer_create(path_bytes)
|
|
176
334
|
_check_ptr(ptr, f"Writer create: {file_path}")
|
|
177
|
-
# The pointer is consumed by finalize, so we don't use ffi.gc here.
|
|
178
|
-
# The writer should be freed via finalize or ztensor_writer_free if finalize fails.
|
|
179
335
|
self._ptr = ptr
|
|
180
336
|
self._finalized = False
|
|
181
337
|
|
|
@@ -183,46 +339,68 @@ class Writer:
|
|
|
183
339
|
return self
|
|
184
340
|
|
|
185
341
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
186
|
-
# Automatically finalize on exit if not already done and no error occurred.
|
|
187
342
|
if self._ptr and not self._finalized:
|
|
188
343
|
if exc_type is None:
|
|
189
344
|
self.finalize()
|
|
190
345
|
else:
|
|
191
|
-
# If an error occurred,
|
|
346
|
+
# If an error occurred, just free the handle without finalizing
|
|
192
347
|
lib.ztensor_writer_free(self._ptr)
|
|
193
348
|
self._ptr = None
|
|
194
349
|
|
|
195
|
-
def add_tensor(self, name: str, tensor
|
|
196
|
-
"""
|
|
197
|
-
|
|
350
|
+
def add_tensor(self, name: str, tensor):
|
|
351
|
+
"""
|
|
352
|
+
Adds a NumPy or PyTorch tensor to the file (zero-copy).
|
|
353
|
+
Supports float16 and bfloat16 types.
|
|
198
354
|
|
|
199
|
-
|
|
200
|
-
|
|
355
|
+
Args:
|
|
356
|
+
name (str): The name of the tensor to add.
|
|
357
|
+
tensor (np.ndarray or torch.Tensor): The tensor data to write.
|
|
358
|
+
"""
|
|
359
|
+
if not self._ptr: raise ZTensorError("Writer is closed or finalized.")
|
|
201
360
|
|
|
202
|
-
|
|
203
|
-
|
|
361
|
+
# --- Polymorphic tensor handling ---
|
|
362
|
+
if isinstance(tensor, np.ndarray):
|
|
363
|
+
tensor = np.ascontiguousarray(tensor)
|
|
364
|
+
shape = tensor.shape
|
|
365
|
+
dtype_str = DTYPE_NP_TO_ZT.get(tensor.dtype)
|
|
366
|
+
data_ptr = ffi.cast("unsigned char*", tensor.ctypes.data)
|
|
367
|
+
nbytes = tensor.nbytes
|
|
368
|
+
|
|
369
|
+
elif TORCH_AVAILABLE and isinstance(tensor, torch.Tensor):
|
|
370
|
+
if tensor.is_cuda:
|
|
371
|
+
raise ZTensorError("Cannot write directly from a CUDA tensor. Copy to CPU first using .cpu().")
|
|
372
|
+
tensor = tensor.contiguous()
|
|
373
|
+
shape = tuple(tensor.shape)
|
|
374
|
+
dtype_str = DTYPE_TORCH_TO_ZT.get(tensor.dtype)
|
|
375
|
+
data_ptr = ffi.cast("unsigned char*", tensor.data_ptr())
|
|
376
|
+
nbytes = tensor.numel() * tensor.element_size()
|
|
377
|
+
|
|
378
|
+
else:
|
|
379
|
+
supported = "np.ndarray" + (" or torch.Tensor" if TORCH_AVAILABLE else "")
|
|
380
|
+
raise TypeError(f"Unsupported tensor type: {type(tensor)}. Must be {supported}.")
|
|
204
381
|
|
|
205
|
-
# The updated DTYPE_NP_TO_ZT will now correctly handle bfloat16 tensors.
|
|
206
|
-
dtype_str = DTYPE_NP_TO_ZT.get(tensor.dtype)
|
|
207
382
|
if not dtype_str:
|
|
208
|
-
|
|
209
|
-
|
|
383
|
+
msg = f"Unsupported dtype: {tensor.dtype}."
|
|
384
|
+
if 'bfloat16' in str(tensor.dtype) and not ML_DTYPES_AVAILABLE:
|
|
385
|
+
msg += " For NumPy bfloat16 support, please install the 'ml_dtypes' package."
|
|
386
|
+
raise ZTensorError(msg)
|
|
210
387
|
|
|
211
|
-
|
|
212
|
-
|
|
388
|
+
name_bytes = name.encode('utf-8')
|
|
389
|
+
shape_array = np.array(shape, dtype=np.uint64)
|
|
390
|
+
shape_ptr = ffi.cast("uint64_t*", shape_array.ctypes.data)
|
|
391
|
+
dtype_bytes = dtype_str.encode('utf-8')
|
|
213
392
|
|
|
214
393
|
status = lib.ztensor_writer_add_tensor(
|
|
215
|
-
self._ptr, name_bytes, shape_ptr, len(
|
|
216
|
-
dtype_bytes, data_ptr,
|
|
394
|
+
self._ptr, name_bytes, shape_ptr, len(shape),
|
|
395
|
+
dtype_bytes, data_ptr, nbytes
|
|
217
396
|
)
|
|
218
397
|
_check_status(status, f"add_tensor: {name}")
|
|
219
398
|
|
|
220
399
|
def finalize(self):
|
|
221
400
|
"""Finalizes the zTensor file, writing the metadata index."""
|
|
222
401
|
if not self._ptr: raise ZTensorError("Writer is already closed or finalized.")
|
|
223
|
-
|
|
224
402
|
status = lib.ztensor_writer_finalize(self._ptr)
|
|
225
|
-
self._ptr = None # The writer
|
|
403
|
+
self._ptr = None # The writer is consumed in Rust
|
|
226
404
|
self._finalized = True
|
|
227
405
|
_check_status(status, "finalize")
|
|
228
406
|
|
ztensor/ztensor/ffi.py
CHANGED
|
@@ -3,8 +3,8 @@ import _cffi_backend
|
|
|
3
3
|
|
|
4
4
|
ffi = _cffi_backend.FFI('ffi',
|
|
5
5
|
_version = 0x2601,
|
|
6
|
-
_types = b'\x00\x00\
|
|
7
|
-
_globals = (b'\xFF\xFF\xFF\x1FALIGNMENT',64,b'\x00\x00\
|
|
8
|
-
_struct_unions = ((b'\x00\x00\x00\
|
|
9
|
-
_typenames = (b'\x00\x00\x00\
|
|
6
|
+
_types = b'\x00\x00\x33\x0D\x00\x00\x4B\x03\x00\x00\x00\x0F\x00\x00\x36\x0D\x00\x00\x4B\x03\x00\x00\x4A\x03\x00\x00\x00\x0F\x00\x00\x39\x0D\x00\x00\x01\x11\x00\x00\x4E\x03\x00\x00\x00\x0F\x00\x00\x39\x0D\x00\x00\x01\x11\x00\x00\x1C\x01\x00\x00\x00\x0F\x00\x00\x04\x0D\x00\x00\x09\x11\x00\x00\x00\x0F\x00\x00\x1B\x0D\x00\x00\x09\x11\x00\x00\x00\x0F\x00\x00\x42\x0D\x00\x00\x05\x11\x00\x00\x00\x0F\x00\x00\x09\x0D\x00\x00\x00\x0F\x00\x00\x4F\x0D\x00\x00\x4C\x03\x00\x00\x00\x0F\x00\x00\x4F\x0D\x00\x00\x1B\x11\x00\x00\x09\x11\x00\x00\x50\x03\x00\x00\x1C\x01\x00\x00\x09\x11\x00\x00\x51\x03\x00\x00\x1C\x01\x00\x00\x00\x0F\x00\x00\x0D\x0D\x00\x00\x05\x11\x00\x00\x00\x0F\x00\x00\x0D\x0D\x00\x00\x01\x11\x00\x00\x00\x0F\x00\x00\x45\x0D\x00\x00\x05\x11\x00\x00\x00\x0F\x00\x00\x50\x0D\x00\x00\x05\x11\x00\x00\x00\x0F\x00\x00\x53\x0D\x00\x00\x48\x03\x00\x00\x00\x0F\x00\x00\x53\x0D\x00\x00\x49\x03\x00\x00\x00\x0F\x00\x00\x53\x0D\x00\x00\x4A\x03\x00\x00\x00\x0F\x00\x00\x53\x0D\x00\x00\x04\x11\x00\x00\x00\x0F\x00\x00\x53\x0D\x00\x00\x1B\x11\x00\x00\x00\x0F\x00\x00\x53\x0D\x00\x00\x4E\x03\x00\x00\x00\x0F\x00\x00\x53\x0D\x00\x00\x50\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\x04\x09\x00\x00\x42\x03\x00\x00\x02\x01\x00\x00\x07\x01\x00\x00\x18\x01\x00\x00\x04\x01\x00\x00\x53\x03\x00\x00\x00\x01',
|
|
7
|
+
_globals = (b'\xFF\xFF\xFF\x1FALIGNMENT',64,b'\x00\x00\x41\x23ztensor_free_string',0,b'\x00\x00\x32\x23ztensor_free_string_array',0,b'\x00\x00\x35\x23ztensor_free_tensor_view',0,b'\x00\x00\x44\x23ztensor_free_u64_array',0,b'\x00\x00\x18\x23ztensor_last_error_message',0,b'\x00\x00\x38\x23ztensor_metadata_free',0,b'\x00\x00\x15\x23ztensor_metadata_get_checksum_str',0,b'\x00\x00\x15\x23ztensor_metadata_get_data_endianness_str',0,b'\x00\x00\x15\x23ztensor_metadata_get_dtype_str',0,b'\x00\x00\x15\x23ztensor_metadata_get_encoding_str',0,b'\x00\x00\x15\x23ztensor_metadata_get_layout_str',0,b'\x00\x00\x15\x23ztensor_metadata_get_name',0,b'\x00\x00\x2F\x23ztensor_metadata_get_offset',0,b'\x00\x00\x2C\x23ztensor_metadata_get_shape_data',0,b'\x00\x00\x26\x23ztensor_metadata_get_shape_len',0,b'\x00\x00\x2F\x23ztensor_metadata_get_size',0,b'\x00\x00\x3B\x23ztensor_reader_free',0,b'\x00\x00\x00\x23ztensor_reader_get_all_tensor_names',0,b'\x00\x00\x0B\x23ztensor_reader_get_metadata_by_index',0,b'\x00\x00\x07\x23ztensor_reader_get_metadata_by_name',0,b'\x00\x00\x29\x23ztensor_reader_get_metadata_count',0,b'\x00\x00\x0F\x23ztensor_reader_open',0,b'\x00\x00\x03\x23ztensor_reader_read_tensor_view',0,b'\x00\x00\x1D\x23ztensor_writer_add_tensor',0,b'\x00\x00\x12\x23ztensor_writer_create',0,b'\x00\x00\x1A\x23ztensor_writer_finalize',0,b'\x00\x00\x3E\x23ztensor_writer_free',0),
|
|
8
|
+
_struct_unions = ((b'\x00\x00\x00\x48\x00\x00\x00\x02CStringArray',b'\x00\x00\x4D\x11strings',b'\x00\x00\x0D\x11len'),(b'\x00\x00\x00\x49\x00\x00\x00\x02CTensorDataView',b'\x00\x00\x23\x11data',b'\x00\x00\x0D\x11len',b'\x00\x00\x52\x11_owner'),(b'\x00\x00\x00\x4A\x00\x00\x00\x10TensorMetadata',),(b'\x00\x00\x00\x4B\x00\x00\x00\x10ZTensorReader_BufReader_File',),(b'\x00\x00\x00\x4C\x00\x00\x00\x10ZTensorWriter_BufWriter_File',)),
|
|
9
|
+
_typenames = (b'\x00\x00\x00\x48CStringArray',b'\x00\x00\x00\x49CTensorDataView',b'\x00\x00\x00\x4ACTensorMetadata',b'\x00\x00\x00\x4BCZTensorReader',b'\x00\x00\x00\x4CCZTensorWriter',b'\x00\x00\x00\x4ATensorMetadata',b'\x00\x00\x00\x4BZTensorReader_BufReader_File',b'\x00\x00\x00\x4CZTensorWriter_BufWriter_File'),
|
|
10
10
|
)
|
ztensor/ztensor/libztensor.so
CHANGED
|
Binary file
|
|
@@ -0,0 +1,9 @@
|
|
|
1
|
+
ztensor-0.1.4.dist-info/METADATA,sha256=ECD-zbYFWEiDg5HYTbi2hthPYh1xKG9y5tBd2VZPSX8,4446
|
|
2
|
+
ztensor-0.1.4.dist-info/WHEEL,sha256=iTWsRz1hnXDQdoNmAgkYGjvp2bRaBb4-srKwYiDluzw,104
|
|
3
|
+
ztensor-0.1.4.dist-info/licenses/LICENSE,sha256=AoeyV1LzTyOz9sbr6uOzk_P0lW963DvhJHnVNVQlI3Y,1063
|
|
4
|
+
ztensor.libs/libgcc_s-5b5488a6.so.1,sha256=HGKUsVmTeNAxEdSy7Ua5Vh_I9FN3RCbPWzvZ7H_TrwE,2749061
|
|
5
|
+
ztensor/__init__.py,sha256=8n2-jVn_Pi95MmfuZeo5K0FKOQwWQLa1hHXM5kBPs2c,15794
|
|
6
|
+
ztensor/ztensor/__init__.py,sha256=sIpB0pJYFX20TdZapIoPMxqMz37wKJyxCkAlTemWDq4,140
|
|
7
|
+
ztensor/ztensor/ffi.py,sha256=xK_OkOYQEMtfTRag-zlNCSLLj2JHB14qMc1cbtqTOB8,3601
|
|
8
|
+
ztensor/ztensor/libztensor.so,sha256=v03q1TASiuFenBWXRXrTQyhh2-UvTVu6thQwQuSFiKg,1636121
|
|
9
|
+
ztensor-0.1.4.dist-info/RECORD,,
|
ztensor-0.1.1.dist-info/RECORD
DELETED
|
@@ -1,9 +0,0 @@
|
|
|
1
|
-
ztensor-0.1.1.dist-info/METADATA,sha256=dpaZnxqi7u30Z3uH641XoVg_DeH5k25UamSBjzseBos,4446
|
|
2
|
-
ztensor-0.1.1.dist-info/WHEEL,sha256=uk6lXffjnM20NjdJS1j5pULyv9smQtZKAlwiwmTYIRE,104
|
|
3
|
-
ztensor-0.1.1.dist-info/licenses/LICENSE,sha256=AoeyV1LzTyOz9sbr6uOzk_P0lW963DvhJHnVNVQlI3Y,1063
|
|
4
|
-
ztensor.libs/libgcc_s-5b5488a6.so.1,sha256=HGKUsVmTeNAxEdSy7Ua5Vh_I9FN3RCbPWzvZ7H_TrwE,2749061
|
|
5
|
-
ztensor/__init__.py,sha256=pzpS0XNcpBb762OPUG7-3raEJxz7P99KX6CGAdh-DS0,9086
|
|
6
|
-
ztensor/ztensor/__init__.py,sha256=sIpB0pJYFX20TdZapIoPMxqMz37wKJyxCkAlTemWDq4,140
|
|
7
|
-
ztensor/ztensor/ffi.py,sha256=J7CG26lx0Xu0IjYR7GWev-BDgGQXseKA4cjkh0IhnLE,2746
|
|
8
|
-
ztensor/ztensor/libztensor.so,sha256=9bgmcP7Tj8o3loKKYgTfK6kzL1wcgxoKCTykU6CSvEI,1611249
|
|
9
|
-
ztensor-0.1.1.dist-info/RECORD,,
|
|
File without changes
|