still an issue with astype

Author: sdpython

onnx_array_api/npx/npx_core_api.py

@@ -123,6 +123,7 @@ def wrapper(*inputs, **kwargs):
             for x in inputs:
                 if isinstance(x, EagerTensor):
                     tensor_class = x.__class__
+                    break
             if tensor_class is None:
                 raise RuntimeError(
                     f"Unable to find an EagerTensor in types "

onnx_array_api/npx/npx_functions.py

@@ -184,8 +184,11 @@ def astype(
     """
     Cast an array.
     """
-    g = a.astype(dtype)
-    return g
+    if isinstance(dtype, Var):
+        raise TypeError(
+            f"dtype is an attribute, it cannot be a Variable of type {type(dtype)}."
+        )
+    return var(a, op="Cast", to=dtype)
 
 
 @npxapi_inline

onnx_array_api/npx/npx_jit_eager.py

@@ -92,6 +92,8 @@ def make_key(*values, **kwargs):
                 res.append(v)
             elif isinstance(v, slice):
                 res.append(("slice", v.start, v.stop, v.step))
+            elif isinstance(v, type):
+                res.append(("type", v.__name__))
             elif isinstance(v, tuple):
                 subkey = []
                 for sk in v:
@@ -139,16 +141,23 @@ def to_jit(self, *values, **kwargs):
         The onnx graph built by the function defines the input
         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 self.output_types is not None:
             constraints.update(self.output_types)
 
-        inputs = [Input(f"x{i}") for i in range(len(values))]
+        inputs = [Input(f"x{i}") for i in range(len(values)) if f"x{i}" in constraints]
+        if len(inputs) < len(values):
+            # An attribute is not named in the numpy API
+            # but is the ONNX definition.
+            raise NotImplementedError()
         var = self.f(*inputs, **kwargs)
 
         onx = var.to_onnx(

onnx_array_api/npx/npx_numpy_tensors.py

@@ -161,6 +161,12 @@ def get_ir_version(cls, ir_version):
         """
         return ir_version
 
+    def const_cast(self, to: Any = None) -> "EagerTensor":
+        """
+        Casts a constant without any ONNX conversion.
+        """
+        return self.__class__(self._tensor.astype(to))
+
     # The class should support whatever Var supports.
     # This part is not yet complete.
 

onnx_array_api/npx/npx_tensors.py

@@ -21,6 +21,15 @@ class EagerTensor(ArrayApi):
     :class:`ArrayApi`.
     """
 
+    def const_cast(self, to: Any = None) -> "EagerTensor":
+        """
+        Casts a constant without any ONNX conversion.
+        """
+        raise NotImplementedError(
+            f"Method 'const_cast' must be overwritten in class "
+            f"{self.__class__.__name__!r}."
+        )
+
     @staticmethod
     def _op_impl(*inputs, method_name=None):
         # avoids circular imports.
@@ -119,7 +128,7 @@ def _generic_method_operator(self, method_name, *args: Any, **kwargs: Any) -> An
         new_args = []
         for a in args:
             if isinstance(a, np.ndarray):
-                new_args.append(self.__class__(a).astype(self.dtype))
+                new_args.append(self.__class__(a).const_cast(self.dtype))
             else:
                 new_args.append(a)