From ee0f4a91a65527c7a41ca665ac9451b3656d3836 Mon Sep 17 00:00:00 2001
From: Yaman Umuroglu <yamanu@xilinx.com>
Date: Tue, 22 Oct 2019 14:20:54 +0100
Subject: [PATCH] [Transform] add infer_shapes

currently wrapper for ONNX shape inference, will change later
---
 src/finn/transformation/infer_shapes.py | 11 +++++++++++
 1 file changed, 11 insertions(+)
 create mode 100644 src/finn/transformation/infer_shapes.py

diff --git a/src/finn/transformation/infer_shapes.py b/src/finn/transformation/infer_shapes.py
new file mode 100644
index 000000000..1f215ec3f
--- /dev/null
+++ b/src/finn/transformation/infer_shapes.py
@@ -0,0 +1,11 @@
+import onnx.shape_inference as si
+
+
+def infer_shapes(model):
+    """Ensure every tensor in the model has a specified shape (ValueInfo)."""
+    # currently just calls ONNX shape inference, but in the future we will
+    # have to handle shape inference for custom ops ourselves
+    model.model = si.infer_shapes(model.model)
+    # single-step operation, no need to call multiple times so return
+    # model_was_changed = false
+    return (model, False)
-- 
GitLab