Update gather documentation to allow index.shape[k] <= input.shape[k]…

… rather than ==.

ghstack-source-id: 10547be51260a46ae09def4f94d09d44773b8562
Pull Request resolved: #41887

torch/_torch_docs.py

@@ -2548,7 +2548,7 @@ def merge_dicts(*dicts):
 
 .. note::  This function is similar to SciPy's `scipy.special.digamma`.
 
-.. note::  From PyTorch 1.8 onwards, the digamma function returns `-Inf` for `0`. 
+.. note::  From PyTorch 1.8 onwards, the digamma function returns `-Inf` for `0`.
            Previously it returned `NaN` for `0`.
 
 Example::
@@ -3127,12 +3127,10 @@ def merge_dicts(*dicts):
     out[i][j][k] = input[i][index[i][j][k]][k]  # if dim == 1
     out[i][j][k] = input[i][j][index[i][j][k]]  # if dim == 2
 
-If :attr:`input` is an n-dimensional tensor with size
-:math:`(x_0, x_1..., x_{i-1}, x_i, x_{i+1}, ..., x_{n-1})`
-and ``dim = i``, then :attr:`index` must be an :math:`n`-dimensional tensor with
-size :math:`(x_0, x_1, ..., x_{i-1}, y, x_{i+1}, ..., x_{n-1})` where :math:`y \geq 1`
-and :attr:`out` will have the same size as :attr:`index`.  Note that ``input``
-and ``index`` do not broadcast against each other.
+:attr:`input` and :attr:`index` must have the same number of dimensions.
+It is also required that ``index.size(d) <= input.size(d)`` for all
+dimensions ``d != dim``.  :attr:`out` will have the same shape as :attr:`index`.
+Note that ``input`` and ``index`` do not broadcast against each other.
 
 Args:
     input (Tensor): the source tensor