fix code

Author: xadupre

_unittests/ut_translate_api/test_translate_classic.py

@@ -406,7 +406,7 @@ def _run(cls, code):
         import onnx.helper
         import onnx.numpy_helper
         import onnx_array_api.translate_api.make_helper
-        import onnx.reference.custom_element_types
+        import ml_dtypes
 
         def from_array_extended(tensor, name=None):
             dt = tensor.dtype
@@ -433,7 +433,7 @@ def from_array_extended(tensor, name=None):
         globs.update(onnx.helper.__dict__)
         globs.update(onnx.numpy_helper.__dict__)
         globs.update(onnx_array_api.translate_api.make_helper.__dict__)
-        globs.update(onnx.reference.custom_element_types.__dict__)
+        globs.update(ml_dtypes.__dict__)
         globs["from_array_extended"] = from_array_extended
         locs = {}
         try:

onnx_array_api/translate_api/base_emitter.py

@@ -129,8 +129,9 @@ def render_attribute_value(self, value: Any) -> Tuple[List[str], str]:
         if value[0].type == AttributeProto.TENSOR:
             repl = {"bool": "bool_", "object": "object_", "str": "str_"}
             sdtype = repl.get(str(v.dtype), str(str(v.dtype)))
+            package = "np" if hasattr(np, sdtype) else "ml_dtypes"
             return [], (
-                f"from_array(np.array({v.tolist()}, dtype=np.{sdtype}), "
+                f"from_array(np.array({v.tolist()}, dtype={package}.{sdtype}), "
                 f"name={value[0].name!r})"
             )
         if isinstance(v, (int, float, list)):

onnx_array_api/translate_api/builder_emitter.py

@@ -1,4 +1,5 @@
 from typing import Any, Dict, List
+import numpy as np
 from onnx import TensorProto
 from onnx.numpy_helper import to_array
 from .base_emitter import BaseEmitter
@@ -135,7 +136,10 @@ def _emit_end_signature(self, **kwargs: Dict[str, Any]) -> List[str]:
             val = to_array(init)
             stype = str(val.dtype).split(".")[-1]
             name = self._clean_result_name(init.name)
-            rows.append(f"    {name} = np.array({val.tolist()}, dtype=np.{stype})")
+            package = "np" if hasattr(np, stype) else "ml_dtypes"
+            rows.append(
+                f"    {name} = np.array({val.tolist()}, dtype={package}.{stype})"
+            )
         return rows
 
     def _emit_begin_return(self, **kwargs: Dict[str, Any]) -> List[str]:

onnx_array_api/translate_api/inner_emitter.py

@@ -1,4 +1,5 @@
 from typing import Any, Dict, List, Optional, Tuple
+import numpy as np
 from onnx import AttributeProto
 from ..annotations import ELEMENT_TYPE_NAME
 from .base_emitter import BaseEmitter
@@ -105,7 +106,7 @@ def _emit_initializer(self, **kwargs: Dict[str, Any]) -> List[str]:
             else:
                 raise NotImplementedError(f"Unexpected dtype={sdtype}.")
         else:
-            sdtype = f"np.{sdtype}"
+            sdtype = f"np.{sdtype}" if hasattr(np, sdtype) else f"ml_dtypes.{sdtype}"
 
         return [
             "initializers.append(",
@@ -233,7 +234,7 @@ def _emit_initializer(self, **kwargs: Dict[str, Any]) -> List[str]:
             else:
                 raise NotImplementedError(f"Unexpected dtype={sdtype}.")
         else:
-            sdtype = f"np.{sdtype}"
+            sdtype = f"np.{sdtype}" if hasattr(np, sdtype) else f"ml_dtypes.{sdtype}"
         if value.size <= 16:
             return [
                 "initializers.append(",

onnx_array_api/translate_api/light_emitter.py

@@ -1,4 +1,5 @@
 from typing import Any, Dict, List
+import numpy as np
 from ..annotations import ELEMENT_TYPE_NAME
 from .base_emitter import BaseEmitter
 
@@ -43,8 +44,9 @@ def _emit_initializer(self, **kwargs: Dict[str, Any]) -> List[str]:
         value = kwargs["value"]
         repl = {"bool": "bool_", "object": "object_", "str": "str_"}
         sdtype = repl.get(str(value.dtype), str(str(value.dtype)))
+        package = "np" if hasattr(np, sdtype) else "ml_dtypes"
         return [
-            f"cst(np.array({value.tolist()}, dtype=np.{sdtype}))",
+            f"cst(np.array({value.tolist()}, dtype={package}.{sdtype}))",
             f"rename({name!r})",
         ]