second step

Author: xadupre

onnx_array_api/npx/npx_functions.py

@@ -1,7 +1,8 @@
 from typing import Any, Optional, Tuple, Union
 
 import numpy as np
-from onnx import FunctionProto, ModelProto, NodeProto
+from onnx import FunctionProto, ModelProto, NodeProto, TensorProto
+from onnx.helper import np_dtype_to_tensor_dtype
 from onnx.numpy_helper import from_array
 
 from .npx_constants import FUNCTION_DOMAIN
@@ -188,7 +189,16 @@ def astype(
         raise TypeError(
             f"dtype is an attribute, it cannot be a Variable of type {type(dtype)}."
         )
-    return var(a, op="Cast", to=dtype)
+    try:
+        to = np_dtype_to_tensor_dtype(dtype)
+    except KeyError:
+        if dtype is int:
+            to = TensorProto.INT64
+        elif dtype is float:
+            to = TensorProto.float64
+        else:
+            raise ValueError(f"Unable to guess tensor type from {dtype}.")
+    return var(a, op="Cast", to=to)
 
 
 @npxapi_inline

onnx_array_api/npx/npx_graph_builder.py

@@ -235,7 +235,13 @@ def make_node(
                 f"Cannot create a node Identity for {len(inputs)} input(s) and "
                 f"{len(outputs)} output(s)."
             )
-        node = make_node(op, inputs, outputs, domain=domain, **new_kwargs)
+        try:
+            node = make_node(op, inputs, outputs, domain=domain, **new_kwargs)
+        except TypeError as e:
+            raise TypeError(
+                f"Unable to create node {op!r}, with inputs={inputs}, "
+                f"outputs={outputs}, domain={domain!r}, new_kwargs={new_kwargs}."
+            ) from e
         for p in protos:
             node.attribute.append(p)
         if attribute_protos is not None:

onnx_array_api/npx/npx_jit_eager.py

@@ -142,30 +142,50 @@ def to_jit(self, *values, **kwargs):
         types and the expected number of dimensions.
         """
         annotations = self.f.__annotations__
-        annot_values = list(annotations.values())
-        constraints = {
-            f"x{i}": v.tensor_type_dims
-            for i, v in enumerate(values)
-            if isinstance(v, (EagerTensor, JitTensor))
-            and (i >= len(annot_values) or issubclass(annot_values[i], TensorType))
-        }
+        if len(annotations) > 0:
+            names = list(annotations.keys())
+            annot_values = list(annotations.values())
+            constraints = {}
+            new_kwargs = {}
+            for i, (v, iname) in enumerate(zip(values, names)):
+                if isinstance(v, (EagerTensor, JitTensor)) and (
+                    i >= len(annot_values) or issubclass(annot_values[i], TensorType)
+                ):
+                    constraints[iname] = v.tensor_type_dims
+                else:
+                    new_kwargs[iname] = v
+        else:
+            names = [f"x{i}" for i in range(len(values))]
+            new_kwargs = {}
+            constraints = {
+                iname: v.tensor_type_dims
+                for i, (v, iname) in enumerate(zip(values, names))
+                if isinstance(v, (EagerTensor, JitTensor))
+            }
 
         if self.output_types is not None:
             constraints.update(self.output_types)
 
-        inputs = [Input(f"x{i}") for i in range(len(values)) if f"x{i}" in constraints]
-        if len(inputs) < len(values):
+        inputs = [
+            Input(iname) for iname, v in zip(names, values) if iname in constraints
+        ]
+        names = [i.name for i in inputs]
+        if len(new_kwargs) > 0:
             # An attribute is not named in the numpy API
             # but is the ONNX definition.
-            raise NotImplementedError()
+            if len(kwargs) == 0:
+                kwargs = new_kwargs
+            else:
+                kwargs = kwargs.copy()
+                kwargs.update(kwargs)
+
         var = self.f(*inputs, **kwargs)
 
         onx = var.to_onnx(
             constraints=constraints,
             target_opsets=self.target_opsets,
             ir_version=self.ir_version,
         )
-        names = [f"x{i}" for i in range(len(values))]
         exe = self.tensor_class.create_function(names, onx)
         return onx, exe
 
@@ -223,7 +243,7 @@ def jit_call(self, *values, **kwargs):
             raise RuntimeError(
                 f"Unable to run function for key={key!r}, "
                 f"types={[type(x) for x in values]}, "
-                f"onnx={self.onxs[key]}."
+                f"kwargs={kwargs}, onnx={self.onxs[key]}."
             ) from e
         return res
 

onnx_array_api/npx/npx_numpy_tensors.py

@@ -37,7 +37,8 @@ def run(self, *inputs: List["NumpyTensor"]) -> List["NumpyTensor"]:
             """
             if len(inputs) != len(self.input_names):
                 raise ValueError(
-                    f"Expected {len(self.input_names)} inputs but got " f"len(inputs)."
+                    f"Expected {len(self.input_names)} inputs but got {len(inputs)}, "
+                    f"self.input_names={self.input_names}, inputs={inputs}."
                 )
             feeds = {}
             for name, inp in zip(self.input_names, inputs):