Feat (quant_tensor): support for dim() and ndim (#947)

Xilinx · Apr 30, 2024 · 33df0c0 · 33df0c0
1 parent cdc1811
commit 33df0c0
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/src/brevitas/quant_tensor/int_quant_tensor.py b/src/brevitas/quant_tensor/int_quant_tensor.py
@@ -215,6 +215,13 @@ def permute(self, *args, **kwargs):
     def size(self, *args, **kwargs):
         return self.value.size(*args, **kwargs)
 
+    @property
+    def ndim(self):
+        return self.value.ndim
+
+    def dim(self):
+        return self.value.dim()
+
     @property
     def shape(self):
         return self.value.shape