Allow subshape selection

README.md

@@ -387,7 +387,7 @@ choice(a)
 
 ### Shaping
 ```
-shape(a)
+shape(a, *dims)
 rank(a)
 length(a) (alias: size)
 isscalar(a)

lab/shaping.py

@@ -63,6 +63,7 @@ def shape(a: Numeric):  # pragma: no cover
 
     Args:
         a (tensor): Tensor.
+        *dims (int, optional): Dimensions to get.
 
     Returns:
         object: Shape of `a`.
@@ -82,6 +83,16 @@ def shape(a: Union[list, tuple]):
     return np.array(a).shape
 
 
+@dispatch
+def shape(a, *dims: Int):
+    a_shape = B.shape(a)
+    subshape = tuple(a_shape[i] for i in dims)
+    if LazyShapes.enabled:
+        return Shape(*subshape)
+    else:
+        return subshape
+
+
 @dispatch
 def rank(a: Union[Numeric, list, tuple]):  # pragma: no cover
     """Get the shape of a tensor.

tests/test_shaping.py

@@ -38,6 +38,7 @@ def test_sizing(f, check_lazy_shapes):
 @pytest.mark.parametrize(
     "x,shape",
     [
+        (1, ()),
         ([], (0,)),
         ([5], (1,)),
         ([[5], [6]], (2, 1)),
@@ -50,6 +51,12 @@ def test_shape(x, shape, check_lazy_shapes):
     assert B.shape(x) == shape
 
 
+def test_subshape(check_lazy_shapes):
+    assert B.shape(B.zeros(2), 0) == (2,)
+    assert B.shape(B.zeros(2, 3, 4), 1) == (3,)
+    assert B.shape(B.zeros(2, 3, 4), 0, 2) == (2, 4)
+
+
 def test_lazy_shape():
     a = B.randn(2, 2)