fix: re-write ivy.Shape.numel to make it traceable

ivy-llc · Jul 10, 2024 · ae17824 · ae17824
1 parent a373675
commit ae17824
Showing 1 changed file with 3 additions and 4 deletions.
diff --git a/ivy/__init__.py b/ivy/__init__.py
@@ -493,10 +493,9 @@ def numel(self):
             raise ivy.utils.exceptions.IvyException(
                 "Cannot calculate the number of elements in a partially known Shape"
             )
-        res = 1
-        for dim in self._shape:
-            res *= dim
-        return res
+        if self.rank > 0:
+            return ivy.prod(self._shape, dtype=ivy.int64).to_scalar()
+        return 1
 
     def __concat__(self, other):
         return self.concatenate(other)