IrrepsData.ndim

e3nn · Mar 26, 2022 · 29ea074 · 29ea074
1 parent 51e6a9b
commit 29ea074
Showing 1 changed file with 4 additions and 0 deletions.
diff --git a/e3nn_jax/_irreps.py b/e3nn_jax/_irreps.py
@@ -764,6 +764,10 @@ def __repr__(self):
     def shape(self):
         return self.contiguous.shape[:-1]
 
+    @property
+    def ndim(self):
+        return len(self.shape)
+
     def reshape(self, shape) -> "IrrepsData":
         list = [
             None if x is None else x.reshape(shape + (mul, ir.dim))