improve shape handling

sdpython · sdpython · Jan 9, 2024 · Jan 9, 2024 · Jan 9, 2024 · Jan 9, 2024
commit 3baf5634a3bb3cf96cfe9ad16814e4e9748a14e4
diff --git a/_unittests/ut_light_api/test_light_api.py b/_unittests/ut_light_api/test_light_api.py
@@ -211,7 +211,7 @@ def test_neg(self):
         self.assertIsInstance(v, Var)
         self.assertEqual(["X"], v.parent.input_names)
         s = str(v)
-        self.assertEqual("X:FLOAT", s)
+        self.assertEqual("X:FLOAT:[]", s)
         onx = start().vin("X").Neg().rename("Y").vout().to_onnx()
         self.assertIsInstance(onx, ModelProto)
         ref = ReferenceEvaluator(onx)
@@ -520,13 +520,13 @@ def test_input_shape(self):
             .bring("X", "W")
             .Conv(pads=[1, 1, 1, 1])
             .rename("Y")
-            .vout()
+            .vout(shape=[])
             .to_onnx()
         )
         i = str(model.graph.input[0]).replace("\n", "").replace(" ", "")
         self.assertNotIn("shape{}", i)
 
 
 if __name__ == "__main__":
-    TestLightApi().test_input_shape()
+    TestLightApi().test_add()
     unittest.main(verbosity=2)
diff --git a/onnx_array_api/annotations.py b/onnx_array_api/annotations.py
@@ -81,9 +81,17 @@ def elem_type_int(elem_type: ELEMENT_TYPE) -> int:
     return np_dtype_to_tensor_dtype(elem_type)
 
 
-def make_shape(shape: TensorShapeProto) -> SHAPE_TYPE:
+def _pick_dim(d, empty_dim):
+    if d.dim_value:
+        return d.dim_value
+    if d.dim_param:
+        return d.dim_param
+    return empty_dim
+
+
+def make_shape(shape: TensorShapeProto, empty_dim: Optional[Any] = None) -> SHAPE_TYPE:
     "Extracts a shape from a tensor type."
-    if hasattr(shape, "dims"):
-        res = [(d.dim_value if d.dim_value else d.dim_param) for d in shape.dims]
+    if hasattr(shape, "dim"):
+        res = [_pick_dim(d, empty_dim=empty_dim) for i, d in enumerate(shape.dim)]
         return tuple(res)
     return None
diff --git a/onnx_array_api/light_api/model.py b/onnx_array_api/light_api/model.py
@@ -180,6 +180,8 @@ def make_output(
         :param elem_type: element type (the input is assumed to be a tensor)
         :param shape: shape
         :return: an instance of ValueInfoProto
+
+        If the checker fails, try `shape=[]`.
         """
         if not self.has_name(name):
             raise ValueError(f"Name {name!r} does not exist.")
@@ -332,7 +334,7 @@ def _fix_name_tensor_input(
     ) -> Union[TensorProto, SparseTensorProto, ValueInfoProto]:
         obj = self._fix_name_tensor(obj)
         shape = make_shape(obj.type.tensor_type.shape)
-        if shape is None:
+        if not shape:
             tensor_type_proto = make_tensor_type_proto(
                 obj.type.tensor_type.elem_type, []
             )
@@ -344,7 +346,7 @@ def _fix_name_tensor_output(
     ) -> Union[TensorProto, SparseTensorProto, ValueInfoProto]:
         obj = self._fix_name_tensor(obj)
         shape = make_shape(obj.type.tensor_type.shape)
-        if shape is None:
+        if not shape:
             tensor_type_proto = make_tensor_type_proto(
                 obj.type.tensor_type.elem_type, []
             )

diff --git a/onnx_array_api/light_api/var.py b/onnx_array_api/light_api/var.py
@@ -318,6 +318,8 @@ def vout(
         :param elem_type: element_type
         :param shape: shape
         :return: instance of :class:`onnx_array_api.light_api.Var`
+
+        If the checker fails, try `shape=[]`.
         """
         output = self.parent.make_output(self.name, elem_type=elem_type, shape=shape)
         return Var(
@@ -461,6 +463,8 @@ def vout(
 
         :param elem_type_shape: list of tuple(element_type, shape)
         :return: instance of :class:`onnx_array_api.light_api.Vars`
+
+        If the checker fails, try `shape=[]`.
         """
         vars = []
         for i, v in enumerate(self.vars_):