DOC Improves shape documentation for *Flatten (#60980)

thomasjpfan · facebook-github-bot · commit 5503a4ac6e07 · 2021-07-06T10:47:11.000-07:00
Summary: Fixes #60841 Pull Request resolved: #60980 Reviewed By: VitalyFedyunin Differential Revision: D29526650 Pulled By: jbschlosser fbshipit-source-id: 2b4b0b84e0652c4cf3e9a48debb3b1bfe4e04b05
diff --git a/torch/nn/modules/flatten.py b/torch/nn/modules/flatten.py
@@ -10,8 +10,10 @@ class Flatten(Module):
     Flattens a contiguous range of dims into a tensor. For use with :class:`~nn.Sequential`.
 
     Shape:
-        - Input: :math:`(N, *dims)`
-        - Output: :math:`(N, \prod *dims)` (for the default case).
+        - Input: :math:`(*, S_{\text{start}},..., S_{i}, ..., S_{\text{end}}, *)`,'
+          where :math:`S_{i}` is the size at dimension :math:`i` and :math:`*` means any
+          number of dimensions including none.
+        - Output: :math:`(*, \prod_{i=\text{start}}^{\text{end}} S_{i}, *)`.
 
     Args:
         start_dim: first dim to flatten (default = 1).
@@ -57,8 +59,10 @@ class Unflatten(Module):
       (tuple of `(name, size)` tuples) for `NamedTensor` input.
 
     Shape:
-        - Input: :math:`(N, *dims)`
-        - Output: :math:`(N, C_{\text{out}}, H_{\text{out}}, W_{\text{out}})`
+        - Input: :math:`(*, S_{\text{dim}}, *)`, where :math:`S_{\text{dim}}` is the size at
+          dimension :attr:`dim` and :math:`*` means any number of dimensions including none.
+        - Output: :math:`(*, U_1, ..., U_n, *)`, where :math:`U` = :attr:`unflattened_size` and
+          :math:`\prod_{i=1}^n U_i = S_{\text{dim}}`.
 
     Args:
         dim (Union[int, str]): Dimension to be unflattened
