better constant

xadupre · xadupre · commit d00d823e57cf · 2023-12-20T15:30:22.000+01:00
diff --git a/onnx_array_api/graph_api/graph_builder.py b/onnx_array_api/graph_api/graph_builder.py
@@ -173,6 +173,9 @@ def _get_tensor_shape(
                     return tuple(att.floats)
                 if att.name == "value_ints":
                     return (len(att.ints),)
+                if att.name == "value":
+                    t = onh.to_array(att.t)
+                    return t.shape
         raise TypeError(
             f"Unexpected or unsupported scenario type {type(proto)}: {proto}."
         )
@@ -190,6 +193,9 @@ def _get_tensor_type(self, proto: Union[NodeProto, TensorProto]) -> int:
                     return TensorProto.FLOAT
                 if att.name == "value_ints":
                     return TensorProto.INT64
+                if att.name == "value":
+                    t = onh.to_array(att.t)
+                    return oh.np_dtype_to_tensor_dtype(t.dtype)
         raise ValueError(f"Unexpected type or value {type(proto)}: {proto}.")
 
     def is_constant(self, name: str) -> bool:
