onnx2fx 0.0.0__py3-none-any.whl

onnx2fx/utils/attributes.py

@@ -0,0 +1,150 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Utilities for parsing ONNX node attributes."""
+
+from typing import TYPE_CHECKING, Any, Callable, Dict, Optional
+
+import onnx
+
+from ..exceptions import UnsupportedDTypeError
+from .dtype import DTYPE_MAP
+
+if TYPE_CHECKING:
+    import torch
+
+
+def get_attribute(
+    node: onnx.NodeProto,
+    name: str,
+    default: Optional[Any] = None,
+    *,
+    tensor_loader: Optional[Callable[[onnx.TensorProto], "torch.Tensor"]] = None,
+) -> Any:
+    """Get a single attribute value from an ONNX node.
+
+    Parameters
+    ----------
+    node : onnx.NodeProto
+        The ONNX node.
+    name : str
+        The attribute name.
+    default : Optional[Any]
+        Default value if attribute is not found.
+
+    Returns
+    -------
+    Any
+        The attribute value, or default if not found.
+    """
+    for attr in node.attribute:
+        if attr.name == name:
+            return _parse_attribute_value(attr, tensor_loader=tensor_loader)
+    return default
+
+
+def get_attributes(node: onnx.NodeProto) -> Dict[str, Any]:
+    """Get all attributes from an ONNX node as a dictionary.
+
+    Parameters
+    ----------
+    node : onnx.NodeProto
+        The ONNX node.
+
+    Returns
+    -------
+    Dict[str, Any]
+        Dictionary mapping attribute names to their values.
+    """
+    return {attr.name: _parse_attribute_value(attr) for attr in node.attribute}
+
+
+def _parse_attribute_value(
+    attr: onnx.AttributeProto,
+    *,
+    tensor_loader: Optional[Callable[[onnx.TensorProto], "torch.Tensor"]] = None,
+) -> Any:
+    """Parse an ONNX attribute into a Python value.
+
+    Parameters
+    ----------
+    attr : onnx.AttributeProto
+        The ONNX attribute.
+
+    Returns
+    -------
+    Any
+        The parsed Python value.
+    """
+    match attr.type:
+        case onnx.AttributeProto.FLOAT:
+            return attr.f
+        case onnx.AttributeProto.INT:
+            return attr.i
+        case onnx.AttributeProto.STRING:
+            return attr.s.decode("utf-8") if isinstance(attr.s, bytes) else attr.s
+        case onnx.AttributeProto.TENSOR:
+            return _parse_tensor(attr.t, tensor_loader=tensor_loader)
+        case onnx.AttributeProto.GRAPH:
+            return attr.g
+        case onnx.AttributeProto.FLOATS:
+            return list(attr.floats)
+        case onnx.AttributeProto.INTS:
+            return list(attr.ints)
+        case onnx.AttributeProto.STRINGS:
+            return [
+                s.decode("utf-8") if isinstance(s, bytes) else s for s in attr.strings
+            ]
+        case onnx.AttributeProto.TENSORS:
+            return [_parse_tensor(t, tensor_loader=tensor_loader) for t in attr.tensors]
+        case onnx.AttributeProto.GRAPHS:
+            return list(attr.graphs)
+        case onnx.AttributeProto.SPARSE_TENSOR:
+            return attr.sparse_tensor
+        case onnx.AttributeProto.SPARSE_TENSORS:
+            return list(attr.sparse_tensors)
+        case onnx.AttributeProto.TYPE_PROTO:
+            return attr.tp
+        case onnx.AttributeProto.TYPE_PROTOS:
+            return list(attr.type_protos)
+        case _:
+            raise ValueError(f"Unsupported attribute type: {attr.type}")
+
+
+def _parse_tensor(
+    tensor: onnx.TensorProto,
+    *,
+    tensor_loader: Optional[Callable[[onnx.TensorProto], "torch.Tensor"]] = None,
+) -> "torch.Tensor":
+    """Convert an ONNX TensorProto to a PyTorch tensor.
+
+    Parameters
+    ----------
+    tensor : onnx.TensorProto
+        The ONNX tensor.
+
+    Returns
+    -------
+    torch.Tensor
+        The converted PyTorch tensor.
+    """
+    if tensor_loader is not None:
+        return tensor_loader(tensor)
+
+    import torch
+    from onnx import numpy_helper
+
+    onnx_dtype = tensor.data_type
+    if DTYPE_MAP.get(onnx_dtype) is None:
+        raise UnsupportedDTypeError(
+            onnx_dtype=onnx_dtype,
+            tensor_name=tensor.name or "<unnamed>",
+            details="attribute tensor dtype not supported",
+        )
+
+    np_array = numpy_helper.to_array(tensor)
+    return torch.from_numpy(np_array.copy())
+
+
+__all__ = [
+    "get_attribute",
+    "get_attributes",
+]

onnx2fx/utils/dtype.py

@@ -0,0 +1,107 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Data type mapping between ONNX and PyTorch."""
+
+from typing import Dict, Optional
+
+import onnx
+import torch
+
+# ONNX TensorProto data type to PyTorch dtype mapping
+DTYPE_MAP: Dict[int, Optional[torch.dtype]] = {
+    onnx.TensorProto.FLOAT: torch.float32,
+    onnx.TensorProto.FLOAT16: torch.float16,
+    onnx.TensorProto.BFLOAT16: torch.bfloat16,
+    onnx.TensorProto.DOUBLE: torch.float64,
+    onnx.TensorProto.INT8: torch.int8,
+    onnx.TensorProto.INT16: torch.int16,
+    onnx.TensorProto.INT32: torch.int32,
+    onnx.TensorProto.INT64: torch.int64,
+    onnx.TensorProto.UINT8: torch.uint8,
+    onnx.TensorProto.UINT16: torch.uint16,
+    onnx.TensorProto.UINT32: torch.uint32,
+    onnx.TensorProto.UINT64: torch.uint64,
+    onnx.TensorProto.BOOL: torch.bool,
+    onnx.TensorProto.COMPLEX64: torch.complex64,
+    onnx.TensorProto.COMPLEX128: torch.complex128,
+    onnx.TensorProto.STRING: None,
+}
+
+# Reverse mapping: PyTorch dtype to ONNX TensorProto data type
+TORCH_TO_ONNX_DTYPE: Dict[torch.dtype, int] = {
+    torch_dtype: onnx_dtype
+    for onnx_dtype, torch_dtype in DTYPE_MAP.items()
+    if torch_dtype is not None
+}
+
+
+def onnx_dtype_to_torch(onnx_dtype: int) -> Optional[torch.dtype]:
+    """Convert ONNX TensorProto data type to PyTorch dtype.
+
+    Parameters
+    ----------
+    onnx_dtype : int
+        ONNX TensorProto data type enum value.
+
+    Returns
+    -------
+    Optional[torch.dtype]
+        Corresponding PyTorch dtype, or None if not supported.
+    """
+    return DTYPE_MAP.get(onnx_dtype)
+
+
+def torch_dtype_to_onnx(torch_dtype: torch.dtype) -> Optional[int]:
+    """Convert PyTorch dtype to ONNX TensorProto data type.
+
+    Parameters
+    ----------
+    torch_dtype : torch.dtype
+        PyTorch dtype.
+
+    Returns
+    -------
+    Optional[int]
+        Corresponding ONNX TensorProto data type, or None if not supported.
+    """
+    return TORCH_TO_ONNX_DTYPE.get(torch_dtype)
+
+
+# ONNX stash_type to PyTorch dtype mapping (used in normalization ops)
+_STASH_TYPE_MAP: Dict[int, torch.dtype] = {
+    1: torch.float32,
+    10: torch.float16,
+    11: torch.float64,
+    16: torch.bfloat16,
+}
+
+
+def stash_type_to_torch_dtype(stash_type: int) -> torch.dtype:
+    """Convert ONNX stash_type attribute to PyTorch dtype.
+
+    The stash_type attribute specifies the floating-point precision
+    for intermediate computations in normalization operations.
+
+    Parameters
+    ----------
+    stash_type : int
+        ONNX stash_type value:
+        - 1: float32
+        - 10: float16
+        - 11: float64
+        - 16: bfloat16
+
+    Returns
+    -------
+    torch.dtype
+        Corresponding PyTorch dtype. Defaults to float32 for unknown values.
+    """
+    return _STASH_TYPE_MAP.get(stash_type, torch.float32)
+
+
+__all__ = [
+    "DTYPE_MAP",
+    "TORCH_TO_ONNX_DTYPE",
+    "onnx_dtype_to_torch",
+    "stash_type_to_torch_dtype",
+    "torch_dtype_to_onnx",
+]

onnx2fx/utils/external_data.py

@@ -0,0 +1,233 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Helpers for ONNX external data tensors."""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+import importlib
+import math
+import os
+from typing import Any, Dict, Iterable, Tuple
+
+import numpy as np
+import onnx
+
+from ..exceptions import ExternalDataError, UnsupportedDTypeError
+from .dtype import DTYPE_MAP
+
+
+def _load_onnx_mapping() -> Any:  # pragma: no cover - version-dependent
+    try:
+        return importlib.import_module("onnx.mapping")
+    except Exception:
+        return importlib.import_module("onnx._mapping")
+
+
+onnx_mapping: Any = _load_onnx_mapping()
+
+
+@dataclass(frozen=True)
+class ExternalDataInfo:
+    """Resolved external data metadata."""
+
+    path: str
+    offset: int
+    length: int
+    shape: Tuple[int, ...]
+    numpy_dtype: np.dtype
+
+
+def _tensor_name(tensor: onnx.TensorProto) -> str:
+    return tensor.name or "<unnamed>"
+
+
+def _get_numpy_dtype(onnx_dtype: int, *, tensor_name: str) -> np.dtype:
+    if DTYPE_MAP.get(onnx_dtype) is None:
+        raise UnsupportedDTypeError(
+            onnx_dtype=onnx_dtype,
+            tensor_name=tensor_name,
+            details="dtype not supported by onnx2fx",
+        )
+    np_type = onnx_mapping.TENSOR_TYPE_TO_NP_TYPE.get(onnx_dtype)
+    if np_type is None:
+        raise UnsupportedDTypeError(
+            onnx_dtype=onnx_dtype,
+            tensor_name=tensor_name,
+            details="dtype has no NumPy mapping",
+        )
+    try:
+        return np.dtype(np_type)
+    except Exception as exc:  # pragma: no cover - defensive
+        raise UnsupportedDTypeError(
+            onnx_dtype=onnx_dtype,
+            tensor_name=tensor_name,
+            details=f"dtype not supported by NumPy memmap ({exc})",
+        ) from exc
+
+
+def _parse_external_data_kv(
+    tensor: onnx.TensorProto, *, tensor_name: str
+) -> Dict[str, str]:
+    data = {entry.key: entry.value for entry in tensor.external_data}
+    if not data:
+        raise ExternalDataError(
+            tensor_name=tensor_name,
+            message="external data metadata is empty",
+        )
+    return data
+
+
+def _require_int_field(
+    data: Dict[str, str],
+    *,
+    field: str,
+    tensor_name: str,
+) -> int:
+    if field not in data:
+        raise ExternalDataError(
+            tensor_name=tensor_name,
+            message=f"missing external data field '{field}'",
+        )
+    try:
+        return int(data[field])
+    except ValueError as exc:
+        raise ExternalDataError(
+            tensor_name=tensor_name,
+            message=f"invalid external data field '{field}': {data[field]}",
+        ) from exc
+
+
+def _resolve_external_path(
+    location: str, base_dir: str | None, *, tensor_name: str
+) -> str:
+    if os.path.isabs(location):
+        return location
+    if not base_dir:
+        raise ExternalDataError(
+            tensor_name=tensor_name,
+            message="base_dir is required for relative external data paths",
+        )
+    return os.path.normpath(os.path.join(base_dir, location))
+
+
+def _expected_nbytes(shape: Iterable[int], np_dtype: np.dtype) -> int:
+    dims = list(shape)
+    if len(dims) == 0:
+        element_count = 1
+    else:
+        element_count = math.prod(dims)
+    return int(element_count) * int(np_dtype.itemsize)
+
+
+def resolve_external_data(
+    tensor: onnx.TensorProto,
+    *,
+    base_dir: str | None,
+    strict: bool = True,
+) -> ExternalDataInfo:
+    tensor_name = _tensor_name(tensor)
+    data = _parse_external_data_kv(tensor, tensor_name=tensor_name)
+    if "location" not in data:
+        raise ExternalDataError(
+            tensor_name=tensor_name,
+            message="missing external data field 'location'",
+        )
+    location = data["location"]
+    offset = _require_int_field(data, field="offset", tensor_name=tensor_name)
+    length = _require_int_field(data, field="length", tensor_name=tensor_name)
+    if offset < 0 or length < 0:
+        raise ExternalDataError(
+            tensor_name=tensor_name,
+            message="external data offset/length must be non-negative",
+        )
+
+    path = _resolve_external_path(location, base_dir, tensor_name=tensor_name)
+    if not os.path.isfile(path):
+        raise ExternalDataError(
+            tensor_name=tensor_name,
+            message=f"external data file not found: {path}",
+        )
+
+    np_dtype = _get_numpy_dtype(tensor.data_type, tensor_name=tensor_name)
+    expected = _expected_nbytes(tensor.dims, np_dtype)
+    file_size = os.path.getsize(path)
+
+    if offset + length > file_size:
+        raise ExternalDataError(
+            tensor_name=tensor_name,
+            message=(
+                "external data range exceeds file size "
+                f"(offset={offset}, length={length}, file_size={file_size})"
+            ),
+        )
+
+    if strict and length != expected:
+        raise ExternalDataError(
+            tensor_name=tensor_name,
+            message=(
+                f"external data length mismatch (length={length}, expected={expected})"
+            ),
+        )
+
+    return ExternalDataInfo(
+        path=path,
+        offset=offset,
+        length=length,
+        shape=tuple(int(dim) for dim in tensor.dims),
+        numpy_dtype=np_dtype,
+    )
+
+
+def iter_all_graphs(model: onnx.ModelProto) -> Iterable[onnx.GraphProto]:
+    """Yield all graphs, including nested subgraphs."""
+    graphs = [model.graph]
+    while graphs:
+        graph = graphs.pop()
+        yield graph
+        for node in graph.node:
+            for attr in node.attribute:
+                if attr.type == onnx.AttributeProto.GRAPH:
+                    graphs.append(attr.g)
+                elif attr.type == onnx.AttributeProto.GRAPHS:
+                    graphs.extend(list(attr.graphs))
+
+
+def validate_external_data_model(
+    model: onnx.ModelProto,
+    *,
+    base_dir: str | None,
+    strict: bool = True,
+) -> None:
+    """Validate external data metadata for all tensors in a model."""
+    for graph in iter_all_graphs(model):
+        for tensor in graph.initializer:
+            if (
+                tensor.data_location == onnx.TensorProto.EXTERNAL
+                or tensor.external_data
+            ):
+                resolve_external_data(tensor, base_dir=base_dir, strict=strict)
+        for node in graph.node:
+            for attr in node.attribute:
+                if attr.type == onnx.AttributeProto.TENSOR:
+                    tensor = attr.t
+                    if (
+                        tensor.data_location == onnx.TensorProto.EXTERNAL
+                        or tensor.external_data
+                    ):
+                        resolve_external_data(tensor, base_dir=base_dir, strict=strict)
+                elif attr.type == onnx.AttributeProto.TENSORS:
+                    for tensor in attr.tensors:
+                        if (
+                            tensor.data_location == onnx.TensorProto.EXTERNAL
+                            or tensor.external_data
+                        ):
+                            resolve_external_data(
+                                tensor, base_dir=base_dir, strict=strict
+                            )
+
+
+__all__ = [
+    "ExternalDataInfo",
+    "resolve_external_data",
+    "validate_external_data_model",
+]

onnx2fx/utils/names.py

@@ -0,0 +1,43 @@
+# SPDX-License-Identifier: Apache-2.0
+"""Utilities for sanitizing ONNX names to valid Python identifiers."""
+
+import keyword
+import re
+
+
+def sanitize_name(name: str) -> str:
+    """Sanitize a name to be a valid Python identifier.
+
+    ONNX tensor names can contain characters that are not valid in Python
+    identifiers (e.g., '.', '/', '-') or may start with digits. This function
+    converts such names to valid Python identifiers.
+
+    Parameters
+    ----------
+    name : str
+        The ONNX tensor name to sanitize.
+
+    Returns
+    -------
+    str
+        A valid Python identifier.
+    """
+    # Replace common invalid characters
+    safe_name = name.replace(".", "_").replace("/", "_").replace("-", "_")
+
+    # Replace any remaining non-alphanumeric characters (except underscore)
+    safe_name = re.sub(r"[^a-zA-Z0-9_]", "_", safe_name)
+
+    # If name starts with a digit, prefix with underscore
+    if safe_name and safe_name[0].isdigit():
+        safe_name = "_" + safe_name
+
+    # Handle empty names
+    if not safe_name:
+        safe_name = "_unnamed"
+
+    # Handle Python keywords
+    if keyword.iskeyword(safe_name):
+        safe_name = safe_name + "_"
+
+    return safe_name