MLBEDSW-3487: Support '<' for tensors
Added __lt__ for Tensor to avoid errors when sorting tensors.
Change-Id: I19bb591ef17aa0d4a3389da411bd8863c2218d55
Signed-off-by: Louis Verhaard <louis.verhaard@arm.com>
diff --git a/ethosu/vela/tensor.py b/ethosu/vela/tensor.py
index 257cb5f..c1443b3 100644
--- a/ethosu/vela/tensor.py
+++ b/ethosu/vela/tensor.py
@@ -21,6 +21,7 @@
from collections import defaultdict
from enum import auto
from functools import lru_cache
+from functools import total_ordering
from typing import Dict
from typing import List
from typing import Optional
@@ -342,6 +343,7 @@
cls.address_map[tens_id][mem_type] = address
+@total_ordering
class Tensor:
__slots__ = (
"shape",
@@ -841,6 +843,9 @@
return (self.dtype.type & BaseType.Int) != 0 and self.quantization.is_valid()
+ def __lt__(self, other: "Tensor") -> bool:
+ return self.equivalence_id < other.equivalence_id
+
def __str__(self):
return "<nng.Tensor '%s' shape=%s dtype=%s>" % (self.name, self.shape, self.dtype)