From 29ea0748e5f7ee38cb35da28d4ea8e800781d0f8 Mon Sep 17 00:00:00 2001
From: Mario Geiger <geiger.mario@gmail.com>
Date: Sat, 26 Mar 2022 11:48:53 -0400
Subject: [PATCH] IrrepsData.ndim

---
 e3nn_jax/_irreps.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/e3nn_jax/_irreps.py b/e3nn_jax/_irreps.py
index 2ff0ad8e..40b8a313 100644
--- 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))