Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.11/site-packages/onnx/numpy_helper.py: 9%
Shortcuts on this page
r m x toggle line displays
j k next/prev highlighted chunk
0 (zero) top of page
1 (one) first highlighted chunk
1# Copyright (c) ONNX Project Contributors
2#
3# SPDX-License-Identifier: Apache-2.0
4from __future__ import annotations
6import math
7import sys
8from collections.abc import Sequence
9from typing import TYPE_CHECKING, Any
11import ml_dtypes
12import numpy as np
13import numpy.typing as npt
15import onnx.external_data_helper
16from onnx import helper
18if TYPE_CHECKING:
19 from collections.abc import Sequence
22def to_float8e8m0(
23 x: np.ndarray,
24 saturate: bool = True,
25 round_mode: str = "up",
26) -> np.ndarray:
27 """Convert a float32 or float64 NumPy array to float8e8m0 representation.
28 Inputs with other dtypes are cast to float32 first.
30 Args:
31 x: Input array to convert.
32 saturate: Whether to saturate at max/min float8e8m0 value.
33 round_mode: "nearest", "up", or "down".
35 Returns:
36 np.ndarray: Array of ml_dtypes.float8_e8m0fnu values.
37 """
38 array = np.asarray(x)
39 if array.dtype.kind != "f" or array.dtype.itemsize not in (4, 8):
40 array = array.astype(np.float32)
42 magnitude = np.abs(array)
43 significand, binary_exponent = np.frexp(magnitude)
44 exponent = binary_exponent.astype(np.int16) - 1
45 exact_power_significand = 0.5
46 nearest_rounding_threshold = 0.75
47 min_exponent = -127
48 max_exponent = 127
50 if round_mode == "nearest":
51 exponent += (significand >= nearest_rounding_threshold).astype(np.int16)
52 elif round_mode == "up":
53 exponent += (significand > exact_power_significand).astype(np.int16)
54 elif round_mode != "down":
55 raise ValueError(f"Unsupported rounding mode: {round_mode}")
57 exponent = np.where(magnitude == 0, min_exponent, exponent)
58 exponent = np.maximum(exponent, min_exponent)
59 if saturate:
60 exponent = np.minimum(exponent, max_exponent)
62 normal_mask = np.isfinite(magnitude) & (exponent <= max_exponent)
63 output = np.where(normal_mask, exponent - min_exponent, 0xFF).astype(np.uint8)
65 return output.view(ml_dtypes.float8_e8m0fnu)
68def _unpack_4bit(
69 data: npt.NDArray[np.uint8], dims: Sequence[int]
70) -> npt.NDArray[np.uint8]:
71 """Convert a packed uint4 array to unpacked uint4 array represented as uint8.
73 Args:
74 data: A numpy array.
75 dims: The dimensions are used to reshape the unpacked buffer.
77 Returns:
78 A numpy array of int8/uint8 reshaped to dims.
79 """
80 result = np.empty([data.size * 2], dtype=data.dtype)
81 array_low = data & np.uint8(0x0F)
82 array_high = data & np.uint8(0xF0)
83 array_high >>= np.uint8(4)
84 result[0::2] = array_low
85 result[1::2] = array_high
86 expected_elements = math.prod(dims)
87 if result.size == expected_elements + 1:
88 # handle single-element padding due to odd number of elements
89 result = result[:-1]
90 if expected_elements > result.size:
91 raise ValueError(
92 f"Packed 4-bit data ({data.size} bytes, {result.size} elements unpacked) "
93 f"is too small for the declared shape {list(dims)} "
94 f"({expected_elements} elements required)."
95 )
96 result.resize(dims, refcheck=False)
97 return result
100def _pack_4bitx2(array: np.ndarray) -> npt.NDArray[np.uint8]:
101 """Convert a numpy array to flatten, packed int4/uint4. Elements must be in the correct range."""
102 # Create a 1D copy
103 array_flat = array.ravel().view(np.uint8).copy()
104 size = array.size
105 odd_sized = size % 2 == 1
106 if odd_sized:
107 array_flat.resize([size + 1], refcheck=False)
108 array_flat &= 0x0F
109 array_flat[1::2] <<= 4
110 return array_flat[0::2] | array_flat[1::2]
113def _unpack_2bit(
114 data: npt.NDArray[np.uint8], dims: Sequence[int]
115) -> npt.NDArray[np.uint8]:
116 """Convert a packed uint2 array to unpacked uint2 array represented as uint8.
118 Args:
119 data: A numpy array.
120 dims: The dimensions are used to reshape the unpacked buffer.
122 Returns:
123 A numpy array of int8/uint8 reshaped to dims.
124 """
125 result = np.empty([data.size * 4], dtype=data.dtype)
126 result[0::4] = data & 0x03
127 result[1::4] = (data >> 2) & 0x03
128 result[2::4] = (data >> 4) & 0x03
129 result[3::4] = (data >> 6) & 0x03
130 expected_elements = math.prod(dims)
131 if result.size > expected_elements:
132 # handle padding due to non multiple of 4 elements
133 result = result[:expected_elements]
134 if expected_elements > result.size:
135 raise ValueError(
136 f"Packed 2-bit data ({data.size} bytes, {result.size} elements unpacked) "
137 f"is too small for the declared shape {list(dims)} "
138 f"({expected_elements} elements required)."
139 )
140 result.resize(dims, refcheck=False)
141 return result
144def _pack_2bitx4(array: np.ndarray) -> npt.NDArray[np.uint8]:
145 """Convert a numpy array to flatten, packed int2/uint2. Elements must be in the correct range."""
146 # Create a 1D copy
147 array_flat = array.ravel().view(np.uint8).copy()
148 size = array.size
149 pad_len = size % 4
150 if pad_len:
151 array_flat.resize([size + (4 - pad_len)], refcheck=False)
152 array_flat &= 0x03
153 array_flat[1::4] <<= 2
154 array_flat[2::4] <<= 4
155 array_flat[3::4] <<= 6
156 return array_flat[0::4] | array_flat[1::4] | array_flat[2::4] | array_flat[3::4]
159def _pack_6bit(values: np.ndarray) -> npt.NDArray[np.uint8]:
160 """Pack a flat uint8 array of 6-bit codes 4-at-a-time into 3 bytes (LSB-first bit order)."""
161 flat = values.astype(np.uint8, copy=False).ravel() & 0x3F
162 n = flat.size
163 packed_size = math.ceil(n * 6 / 8)
164 pad = -n % 4
165 if pad:
166 flat = np.concatenate([flat, np.zeros(pad, dtype=np.uint8)])
167 v0, v1, v2, v3 = flat[0::4], flat[1::4], flat[2::4], flat[3::4]
168 packed = np.empty((v0.size, 3), dtype=np.uint8)
169 packed[:, 0] = v0 | ((v1 & 0x03) << 6)
170 packed[:, 1] = (v1 >> 2) | ((v2 & 0x0F) << 4)
171 packed[:, 2] = (v2 >> 4) | (v3 << 2)
172 return packed.reshape(-1)[:packed_size]
175def _unpack_6bit(data: np.ndarray, dims: Sequence[int]) -> npt.NDArray[np.uint8]:
176 """Unpack a 6-bit packed buffer (see _pack_6bit) back into a uint8 array of given dims."""
177 original_size = math.prod(dims)
178 num_groups = -(-original_size // 4) # ceil division
179 needed_bytes = num_groups * 3
180 # _pack_6bit trims its output to this minimal size: trailing bytes beyond
181 # it would only ever encode a partial final group's zero-padding bits.
182 min_bytes = -(-original_size * 6 // 8) # ceil division
183 data = data.astype(np.uint8, copy=False)
184 if data.size < min_bytes:
185 raise ValueError(
186 f"Packed 6-bit data ({data.size} bytes) is too small for the declared "
187 f"shape {list(dims)} ({min_bytes} bytes required)."
188 )
189 # Bulk-unpack whole 3-byte groups via strided views (no copy of `data`);
190 # only the possibly-incomplete final group (at most 2 missing bytes) needs
191 # padding, so pad just that instead of copying the whole buffer.
192 bulk_bytes = min(data.size, needed_bytes) // 3 * 3
193 bulk_groups = bulk_bytes // 3
194 unpacked = np.empty((num_groups, 4), dtype=np.uint8)
195 b0, b1, b2 = data[0:bulk_bytes:3], data[1:bulk_bytes:3], data[2:bulk_bytes:3]
196 unpacked[:bulk_groups, 0] = b0 & 0x3F
197 unpacked[:bulk_groups, 1] = ((b0 >> 6) & 0x03) | ((b1 & 0x0F) << 2)
198 unpacked[:bulk_groups, 2] = ((b1 >> 4) & 0x0F) | ((b2 & 0x03) << 4)
199 unpacked[:bulk_groups, 3] = (b2 >> 2) & 0x3F
200 if bulk_groups < num_groups:
201 tail = np.zeros(3, dtype=np.uint8)
202 rest = data[bulk_bytes:]
203 tail[: rest.size] = rest
204 t0, t1, t2 = tail
205 unpacked[bulk_groups, 0] = t0 & 0x3F
206 unpacked[bulk_groups, 1] = ((t0 >> 6) & 0x03) | ((t1 & 0x0F) << 2)
207 unpacked[bulk_groups, 2] = ((t1 >> 4) & 0x0F) | ((t2 & 0x03) << 4)
208 unpacked[bulk_groups, 3] = (t2 >> 2) & 0x3F
209 return unpacked.reshape(-1)[:original_size].reshape(dims)
212def _reshape_or_raise(
213 array: np.ndarray, dims: Sequence[int], tensor: onnx.TensorProto
214) -> np.ndarray:
215 """Reshape ``array`` to ``dims``, raising a more actionable error on mismatch.
217 A size mismatch here often means the bytes being parsed are not actually a
218 TensorProto: protobuf does not encode message-type information, so bytes
219 from a different message type (e.g. OptionalProto, SequenceProto) can be
220 silently misparsed as a TensorProto instead of failing to parse.
221 """
222 try:
223 return array.reshape(dims)
224 except ValueError as e:
225 raise ValueError(
226 f"{e} Failed to reshape tensor '{tensor.name}'. This can happen "
227 "when the parsed bytes are not actually a TensorProto; protobuf "
228 "does not encode message-type information, so a different message "
229 "type (e.g. OptionalProto, SequenceProto, MapProto) can be "
230 "silently misparsed as one."
231 ) from e
234def to_array(tensor: onnx.TensorProto, base_dir: str = "") -> np.ndarray: # noqa: PLR0911
235 """Converts a tensor def object to a numpy array.
237 This function uses ml_dtypes if the dtype is not a native numpy dtype.
239 Args:
240 tensor: a TensorProto object.
241 base_dir: if external tensor exists, base_dir can help to find the path to it
243 Returns:
244 arr: the converted array.
245 """
246 if tensor.HasField("segment"):
247 raise ValueError("Currently not supporting loading segments.")
248 if tensor.data_type == onnx.TensorProto.UNDEFINED:
249 raise TypeError("The element type in the input tensor is UNDEFINED.")
251 tensor_dtype = tensor.data_type
252 np_dtype = helper.tensor_dtype_to_np_dtype(tensor_dtype)
253 storage_np_dtype = helper.tensor_dtype_to_np_dtype(
254 helper.tensor_dtype_to_storage_tensor_dtype(tensor_dtype)
255 )
256 storage_field = helper.tensor_dtype_to_field(tensor_dtype)
257 dims = tensor.dims
259 if tensor.data_type == onnx.TensorProto.STRING:
260 utf8_strings = getattr(tensor, storage_field)
261 ss = [s.decode("utf-8") for s in utf8_strings]
262 return _reshape_or_raise(np.asarray(ss).astype(np_dtype), dims, tensor)
264 # Load raw data from external tensor if it exists, without mutating the
265 # input tensor: writing the data back into `tensor.raw_data` would leave
266 # the tensor's data_location/external_data fields stale, corrupting a
267 # subsequent save of the model.
268 external_raw_data = None
269 if onnx.external_data_helper.uses_external_data(tensor):
270 external_raw_data = onnx.external_data_helper._read_external_data_bytes(
271 tensor, base_dir
272 )
274 if external_raw_data is not None or tensor.HasField("raw_data"):
275 # Raw_bytes support: using frombuffer.
276 raw_data = (
277 external_raw_data if external_raw_data is not None else tensor.raw_data
278 )
279 if sys.byteorder == "big":
280 # Convert endian from little to big
281 raw_data = np.frombuffer(raw_data, dtype=np_dtype).byteswap().tobytes()
283 if tensor_dtype in {
284 onnx.TensorProto.INT4,
285 onnx.TensorProto.UINT4,
286 onnx.TensorProto.FLOAT4E2M1,
287 }:
288 data = np.frombuffer(raw_data, dtype=np.uint8)
289 return _unpack_4bit(data, dims).view(np_dtype)
291 if tensor_dtype in {
292 onnx.TensorProto.UINT2,
293 onnx.TensorProto.INT2,
294 }:
295 data = np.frombuffer(raw_data, dtype=np.uint8)
296 return _unpack_2bit(data, dims).view(np_dtype)
298 if tensor_dtype in {onnx.TensorProto.FLOAT6E2M3, onnx.TensorProto.FLOAT6E3M2}:
299 # raw_data is always the packed 6-bit stream (matching UINT4/UINT2's
300 # single-format convention below). A byte-length-based sniff between
301 # "packed" and "one byte per element" would be ambiguous for small
302 # tensors -- e.g. 3 elements pack to ceil(3*6/8) = 3 bytes too.
303 data = np.frombuffer(raw_data, dtype=np.uint8)
304 unpacked = _unpack_6bit(data, dims)
305 return unpacked.view(np_dtype)
307 return _reshape_or_raise(np.frombuffer(raw_data, dtype=np_dtype), dims, tensor)
309 if tensor_dtype in {
310 onnx.TensorProto.BFLOAT16,
311 onnx.TensorProto.FLOAT16,
312 onnx.TensorProto.INT16,
313 onnx.TensorProto.UINT16,
314 }:
315 return _reshape_or_raise(
316 np.array(tensor.int32_data, dtype=np.int32)
317 .view(np.uint32)
318 .astype(np.uint16),
319 dims,
320 tensor,
321 ).view(np_dtype)
323 if tensor_dtype in {
324 onnx.TensorProto.FLOAT8E4M3FN,
325 onnx.TensorProto.FLOAT8E4M3FNUZ,
326 onnx.TensorProto.FLOAT8E5M2,
327 onnx.TensorProto.FLOAT8E5M2FNUZ,
328 onnx.TensorProto.FLOAT8E8M0,
329 onnx.TensorProto.BOOL,
330 }:
331 return _reshape_or_raise(
332 np.array(tensor.int32_data, dtype=np.int32)
333 .view(np.uint32)
334 .astype(np.uint8)
335 .view(np_dtype),
336 dims,
337 tensor,
338 )
340 if tensor_dtype in {
341 onnx.TensorProto.FLOAT6E2M3,
342 onnx.TensorProto.FLOAT6E3M2,
343 }:
344 # Only the low 6 bits of each int32_data entry are meaningful; mask
345 # before reinterpreting, matching _pack_6bit's masking of the packed
346 # raw_data representation.
347 return _reshape_or_raise(
348 (
349 np.array(tensor.int32_data, dtype=np.int32)
350 .view(np.uint32)
351 .astype(np.uint8)
352 & 0x3F
353 ).view(np_dtype),
354 dims,
355 tensor,
356 )
358 if tensor_dtype in {
359 onnx.TensorProto.UINT4,
360 onnx.TensorProto.INT4,
361 onnx.TensorProto.FLOAT4E2M1,
362 }:
363 data = (
364 np.array(tensor.int32_data, dtype=np.int32).view(np.uint32).astype(np.uint8)
365 )
366 return _unpack_4bit(data, dims).view(np_dtype)
368 if tensor_dtype in {
369 onnx.TensorProto.UINT2,
370 onnx.TensorProto.INT2,
371 }:
372 data = (
373 np.array(tensor.int32_data, dtype=np.int32).view(np.uint32).astype(np.uint8)
374 )
375 return _unpack_2bit(data, dims).view(np_dtype)
377 data = getattr(tensor, storage_field)
378 if tensor_dtype in (onnx.TensorProto.COMPLEX64, onnx.TensorProto.COMPLEX128):
379 return _reshape_or_raise(
380 np.array(data, dtype=storage_np_dtype).view(dtype=np_dtype), dims, tensor
381 )
383 return _reshape_or_raise(
384 np.asarray(data, dtype=storage_np_dtype).astype(np_dtype), dims, tensor
385 )
388def tobytes_little_endian(array: np.ndarray) -> bytes:
389 """Converts an array into bytes in little endian byte order.
391 Args:
392 array: a numpy array.
394 Returns:
395 bytes: Byte representation of passed array in little endian byte order.
397 .. versionadded:: 1.20
398 """
399 if array.dtype.byteorder == ">" or (
400 sys.byteorder == "big" and array.dtype.byteorder == "="
401 ):
402 # Ensure that the bytes will be in little-endian byte-order.
403 array = array.astype(array.dtype.newbyteorder("<"))
405 return array.tobytes()
408def from_array(array: np.ndarray, /, name: str | None = None) -> onnx.TensorProto:
409 """Converts an array into a TensorProto including
411 Args:
412 array: a numpy array.
413 name: (optional) the name of the tensor.
415 Returns:
416 TensorProto: the converted tensor def.
417 """
418 tensor = onnx.TensorProto()
419 tensor.dims.extend(array.shape)
420 if name:
421 tensor.name = name
422 if array.dtype == object or np.issubdtype(array.dtype, np.str_):
423 # Special care for strings.
424 tensor.data_type = onnx.TensorProto.STRING
425 # TODO: Introduce full string support.
426 # We flatten the array in case there are n-D arrays are specified
427 # If you want more complex shapes then follow the below instructions.
428 # Unlike other types where the shape is automatically inferred from
429 # nested arrays of values, the only reliable way now to feed strings
430 # is to put them into a flat array then specify type astype(object)
431 # (otherwise all strings may have different types depending on their length)
432 # and then specify shape .reshape([x, y, z])
433 flat_array = array.flatten()
434 for e in flat_array:
435 if isinstance(e, str):
436 tensor.string_data.append(e.encode("utf-8"))
437 elif isinstance(e, bytes):
438 tensor.string_data.append(e)
439 else:
440 raise NotImplementedError(
441 f"Unrecognized object in the object array, expect a string, or array of bytes: {type(e)}"
442 )
443 return tensor
445 dtype = helper.np_dtype_to_tensor_dtype(array.dtype)
446 if dtype in {
447 onnx.TensorProto.INT4,
448 onnx.TensorProto.UINT4,
449 onnx.TensorProto.FLOAT4E2M1,
450 }:
451 # Pack the array into int4
452 array = _pack_4bitx2(array)
454 if dtype in {
455 onnx.TensorProto.UINT2,
456 onnx.TensorProto.INT2,
457 }:
458 # Pack the array into int2
459 array = _pack_2bitx4(array)
461 if dtype in {
462 onnx.TensorProto.FLOAT6E2M3,
463 onnx.TensorProto.FLOAT6E3M2,
464 }:
465 # Pack the array into 6-bit codes
466 array = _pack_6bit(array.view(np.uint8))
468 tensor.raw_data = tobytes_little_endian(array)
469 tensor.data_type = dtype # type: ignore[assignment]
470 return tensor
473def to_list(sequence: onnx.SequenceProto) -> list[Any]:
474 """Converts a sequence def to a Python list.
476 Args:
477 sequence: a SequenceProto object.
479 Returns:
480 list: the converted list.
481 """
482 elem_type = sequence.elem_type
483 if elem_type == onnx.SequenceProto.TENSOR:
484 return [to_array(v) for v in sequence.tensor_values]
485 if elem_type == onnx.SequenceProto.SPARSE_TENSOR:
486 return [to_array(v) for v in sequence.sparse_tensor_values] # type: ignore[arg-type]
487 if elem_type == onnx.SequenceProto.SEQUENCE:
488 return [to_list(v) for v in sequence.sequence_values]
489 if elem_type == onnx.SequenceProto.MAP:
490 return [to_dict(v) for v in sequence.map_values]
491 raise TypeError("The element type in the input sequence is not supported.")
494def from_list(
495 lst: list[Any], name: str | None = None, dtype: int | None = None
496) -> onnx.SequenceProto:
497 """Converts a list into a sequence def.
499 Args:
500 lst: a Python list
501 name: (optional) the name of the sequence.
502 dtype: (optional) type of element in the input list, used for specifying
503 sequence values when converting an empty list.
505 Returns:
506 SequenceProto: the converted sequence def.
507 """
508 sequence = onnx.SequenceProto()
509 if name:
510 sequence.name = name
512 if dtype is not None:
513 elem_type = dtype
514 elif len(lst) > 0:
515 first_elem = lst[0]
516 if isinstance(first_elem, dict):
517 elem_type = onnx.SequenceProto.MAP
518 elif isinstance(first_elem, list):
519 elem_type = onnx.SequenceProto.SEQUENCE
520 else:
521 elem_type = onnx.SequenceProto.TENSOR
522 else:
523 # if empty input list and no dtype specified
524 # choose sequence of tensors on default
525 elem_type = onnx.SequenceProto.TENSOR
526 sequence.elem_type = elem_type
528 if (len(lst) > 0) and not all(isinstance(elem, type(lst[0])) for elem in lst):
529 raise TypeError(
530 "The element type in the input list is not the same "
531 "for all elements and therefore is not supported as a sequence."
532 )
534 if elem_type == onnx.SequenceProto.TENSOR:
535 for tensor in lst:
536 sequence.tensor_values.extend([from_array(np.asarray(tensor))])
537 elif elem_type == onnx.SequenceProto.SEQUENCE:
538 for seq in lst:
539 sequence.sequence_values.extend([from_list(seq)])
540 elif elem_type == onnx.SequenceProto.MAP:
541 for mapping in lst:
542 sequence.map_values.extend([from_dict(mapping)])
543 else:
544 raise TypeError(
545 "The element type in the input list is not a tensor, "
546 "sequence, or map and is not supported."
547 )
548 return sequence
551def to_dict(map_proto: onnx.MapProto) -> dict[Any, Any]:
552 """Converts a map def to a Python dictionary.
554 Args:
555 map_proto: a MapProto object.
557 Returns:
558 The converted dictionary.
559 """
560 key_list: list[Any] = []
561 if map_proto.key_type == onnx.TensorProto.STRING:
562 key_list = list(map_proto.string_keys)
563 else:
564 key_list = list(map_proto.keys)
566 value_list = to_list(map_proto.values)
567 if len(key_list) != len(value_list):
568 raise IndexError(
569 f"Length of keys and values for MapProto (map name: {map_proto.name}) are not the same."
570 )
571 return dict(zip(key_list, value_list, strict=False))
574def _map_key_dtype(key: Any) -> np.dtype:
575 # np.result_type interprets a str argument as a dtype name, so string
576 # keys (of any length) are all mapped to the dtype used for STRING.
577 if isinstance(key, (str, bytes)):
578 return np.dtype(object)
579 return np.result_type(key)
582def from_dict(dict_: dict[Any, Any], name: str | None = None) -> onnx.MapProto:
583 """Converts a Python dictionary into a map def.
585 Args:
586 dict_: Python dictionary
587 name: (optional) the name of the map.
589 Returns:
590 MapProto: the converted map def.
591 """
592 map_proto = onnx.MapProto()
593 if name:
594 map_proto.name = name
595 if not dict_:
596 raise ValueError("Cannot convert an empty dictionary to MapProto.")
597 keys = list(dict_)
598 raw_key_type = _map_key_dtype(keys[0])
599 key_type = helper.np_dtype_to_tensor_dtype(raw_key_type)
601 valid_key_int_types = {
602 onnx.TensorProto.INT8,
603 onnx.TensorProto.INT16,
604 onnx.TensorProto.INT32,
605 onnx.TensorProto.INT64,
606 onnx.TensorProto.UINT8,
607 onnx.TensorProto.UINT16,
608 onnx.TensorProto.UINT32,
609 onnx.TensorProto.UINT64,
610 }
612 if not (all(_map_key_dtype(key) == raw_key_type for key in keys)):
613 raise TypeError(
614 "The key type in the input dictionary is not the same "
615 "for all keys and therefore is not valid as a map."
616 )
618 values = list(dict_.values())
619 raw_value_type = np.result_type(values[0])
620 if not all(np.result_type(val) == raw_value_type for val in values):
621 raise TypeError(
622 "The value type in the input dictionary is not the same "
623 "for all values and therefore is not valid as a map."
624 )
626 value_seq = from_list(values)
628 map_proto.key_type = key_type # type: ignore[assignment]
629 if key_type == onnx.TensorProto.STRING:
630 string_keys = [
631 key.encode("utf-8") if isinstance(key, str) else key for key in keys
632 ]
633 if len(set(string_keys)) != len(string_keys):
634 raise ValueError(
635 "The keys in the input dictionary are not unique after "
636 "encoding str keys as UTF-8 bytes and therefore are not "
637 "valid as a map."
638 )
639 map_proto.string_keys.extend(string_keys)
640 elif key_type in valid_key_int_types:
641 map_proto.keys.extend(keys)
642 else:
643 raise TypeError(f"Unsupported map key type: {key_type}")
644 map_proto.values.CopyFrom(value_seq)
645 return map_proto
648def to_optional(optional: onnx.OptionalProto) -> Any | None:
649 """Converts an optional def to a Python optional.
651 Args:
652 optional: an OptionalProto object.
654 Returns:
655 opt: the converted optional.
656 """
657 elem_type = optional.elem_type
658 if elem_type == onnx.OptionalProto.UNDEFINED:
659 return None
660 if elem_type == onnx.OptionalProto.TENSOR:
661 return (
662 to_array(optional.tensor_value)
663 if optional.HasField("tensor_value")
664 else None
665 )
666 if elem_type == onnx.OptionalProto.SPARSE_TENSOR:
667 return (
668 to_array(optional.sparse_tensor_value) # type: ignore[arg-type]
669 if optional.HasField("sparse_tensor_value")
670 else None
671 )
672 if elem_type == onnx.OptionalProto.SEQUENCE:
673 return (
674 to_list(optional.sequence_value)
675 if optional.HasField("sequence_value")
676 else None
677 )
678 if elem_type == onnx.OptionalProto.MAP:
679 return to_dict(optional.map_value) if optional.HasField("map_value") else None
680 if elem_type == onnx.OptionalProto.OPTIONAL:
681 return (
682 to_optional(optional.optional_value)
683 if optional.HasField("optional_value")
684 else None
685 )
686 raise TypeError("The element type in the input optional is not supported.")
689def from_optional(
690 opt: Any | None, name: str | None = None, dtype: int | None = None
691) -> onnx.OptionalProto:
692 """Converts an optional value into a Optional def.
694 Args:
695 opt: a Python optional
696 name: (optional) the name of the optional.
697 dtype: (optional) type of element in the input, used for specifying
698 optional values when converting empty none. dtype must
699 be a valid OptionalProto.DataType value
701 Returns:
702 optional: the converted optional def.
703 """
704 # TODO: create a map and replace conditional branches
705 optional = onnx.OptionalProto()
706 if name:
707 optional.name = name
709 if dtype is not None:
710 # dtype must be a valid onnx.OptionalProto.DataType
711 if dtype not in onnx.OptionalProto.DataType.values():
712 raise TypeError(f"{dtype} must be a valid OptionalProto.DataType.")
713 elem_type = dtype
714 elif isinstance(opt, dict):
715 elem_type = onnx.OptionalProto.MAP
716 elif isinstance(opt, list):
717 elem_type = onnx.OptionalProto.SEQUENCE
718 elif opt is None:
719 elem_type = onnx.OptionalProto.UNDEFINED
720 else:
721 elem_type = onnx.OptionalProto.TENSOR
723 optional.elem_type = elem_type
725 if opt is not None:
726 if elem_type == onnx.OptionalProto.TENSOR:
727 optional.tensor_value.CopyFrom(from_array(opt))
728 elif elem_type == onnx.OptionalProto.SEQUENCE:
729 optional.sequence_value.CopyFrom(from_list(opt))
730 elif elem_type == onnx.OptionalProto.MAP:
731 optional.map_value.CopyFrom(from_dict(opt))
732 else:
733 raise TypeError(
734 "The element type in the input is not a tensor, "
735 "sequence, or map and is not supported."
736 )
737 return optional
740def create_random_int(
741 input_shape: tuple[int], dtype: np.dtype, seed: int = 1
742) -> np.ndarray:
743 """Create random integer array for backend/test/case/node.
745 Args:
746 input_shape: The shape for the returned integer array.
747 dtype: The NumPy data type for the returned integer array.
748 seed: The seed for np.random.
750 Returns:
751 np.ndarray: Random integer array.
752 """
753 np.random.seed(seed)
754 if dtype in (
755 np.uint8,
756 np.uint16,
757 np.uint32,
758 np.uint64,
759 np.int8,
760 np.int16,
761 np.int32,
762 np.int64,
763 ):
764 # the range of np.random.randint is int32; set a fixed boundary if overflow
765 end = min(np.iinfo(dtype).max, np.iinfo(np.int32).max)
766 start = max(np.iinfo(dtype).min, np.iinfo(np.int32).min)
767 return np.random.randint(start, end, size=input_shape).astype(dtype)
768 raise TypeError(f"{dtype} is not supported by create_random_int.")
771def saturate_cast(x: np.ndarray, dtype: np.dtype) -> np.ndarray:
772 """Saturate cast for numeric types.
774 This function ensures that values outside the representable range
775 of the target dtype are clamped to the maximum or minimum representable
776 value of that dtype.
777 """
778 if np.issubdtype(dtype, np.integer) or dtype in (
779 ml_dtypes.int4,
780 ml_dtypes.uint4,
781 ml_dtypes.int2,
782 ml_dtypes.uint2,
783 ):
784 info = ml_dtypes.iinfo(dtype)
785 x = np.round(x)
786 else:
787 info = ml_dtypes.finfo(dtype) # type: ignore[assignment]
789 return np.clip(x, info.min, info.max).astype(dtype) # type: ignore[no-any-return]