Revert "Revert "MLBEDSW-3645 4D class for op ifm/ofm shapes""
This reverts commit df0a5905177f3a1b836076bc3f9f39b2e86f1794.
Reason for revert: <INSERT REASONING HERE>
Change-Id: I891c66fb29db9d25e942947e8d1c29a10610de51
diff --git a/ethosu/vela/nn_graph.py b/ethosu/vela/nn_graph.py
index 6792517..d2c848a 100644
--- a/ethosu/vela/nn_graph.py
+++ b/ethosu/vela/nn_graph.py
@@ -21,8 +21,10 @@
# Subgraph - Holds a neural network subgraph, pointing at Tensors, Operations, Passes, and CascadedPasses.
# Graph - A full neural network graph with one or more Subgraphs.
import enum
+from typing import List
from .operation import Op
+from .shape4d import Shape4D
class PassPlacement(enum.Enum):
@@ -58,8 +60,8 @@
self.name = name
self.cascade = None
self.placement = placement
- self.ifm_shapes = []
- self.ofm_shapes = []
+ self.ifm_shapes: List[Shape4D] = []
+ self.ofm_shapes: List[Shape4D] = []
# TODO: rename is_element_wise because it is not the same as an ElementWise operator. It is used by the tensor
# allocation and requires that the OFM and IFM has the exact same address. Essentially complete overlap.