add ndim property

src/totypes/types.py

@@ -80,9 +80,14 @@ def __post_init__(self) -> None:
 
     @property
     def shape(self) -> Tuple[int, ...]:
-        """Returns the shape of the array."""
+        """Return the shape of the array."""
         return jnp.shape(self.array)  # type: ignore[no-any-return]
 
+    @property
+    def ndim(self) -> int:
+        """Return the number of dimensions of the array."""
+        return int(jnp.ndim(self.array))
+
 
 def _flatten_bounded_array(
     bounded_array: BoundedArray,
@@ -252,9 +257,14 @@ def __post_init__(self) -> None:
 
     @property
     def shape(self) -> Tuple[int, ...]:
-        """Returns the shape of the array."""
+        """Return the shape of the array."""
         return jnp.shape(self.array)  # type: ignore[no-any-return]
 
+    @property
+    def ndim(self) -> int:
+        """Return the number of dimensions of the array."""
+        return int(jnp.ndim(self.array))
+
 
 def _flatten_density_2d(
     density: Density2DArray,