improved docs examples readability

ml-explore · Feb 9, 2024 · 81a5ae1 · 81a5ae1
1 parent 69f0643
commit 81a5ae1
Showing 1 changed file with 10 additions and 34 deletions.
diff --git a/python/mlx/nn/layers/upsample.py b/python/mlx/nn/layers/upsample.py
@@ -84,41 +84,17 @@ class Upsample2d(Module):
                 [[3],
                  [4]]]], dtype=int32)
         >>> n = nn.Upsample2d(scale=2, mode='nearest')
-        >>> n(x)
-        array([[[[1],
-                 [1],
-                 [2],
-                 [2]],
-                [[1],
-                 [1],
-                 [2],
-                 [2]],
-                [[3],
-                 [3],
-                 [4],
-                 [4]],
-                [[3],
-                 [3],
-                 [4],
-                 [4]]]], dtype=int32)
+        >>> n(x).squeeze()
+        array([[1, 1, 2, 2],
+               [1, 1, 2, 2],
+               [3, 3, 4, 4],
+               [3, 3, 4, 4]], dtype=int32)
         >>> b = nn.Upsample2d(scale=2, mode='bilinear')
-        >>> b(x)
-        array([[[[1],
-                 [1.33333],
-                 [1.66667],
-                 [2]],
-                [[1.66667],
-                 [2],
-                 [2.33333],
-                 [2.66667]],
-                [[2.33333],
-                 [2.66667],
-                 [3],
-                 [3.33333]],
-                [[3],
-                 [3.33333],
-                 [3.66667],
-                 [4]]]], dtype=float32)
+        >>> b(x).squeeze()
+        array([[1, 1.33333, 1.66667, 2],
+               [1.66667, 2, 2.33333, 2.66667],
+               [2.33333, 2.66667, 3, 3.33333],
+               [3, 3.33333, 3.66667, 4]], dtype=float32)
     """
 
     def __init__(