docs: improve torch.matrix_exp() (pytorch#55626)

Summary:
Add a signature and make the mathematical expression related to the signature

Fixes pytorch#55599

Pull Request resolved: pytorch#55626

Reviewed By: ngimel

Differential Revision: D27699518

Pulled By: mruberry

fbshipit-source-id: e61d76e99eb8fc36114c1c2ee90990740d78beea

torch/_torch_docs.py

@@ -5000,13 +5000,14 @@ def merge_dicts(*dicts):
 Alias for :func:`torch.linalg.matrix_power`
 """.format(**common_args))
 
-add_docstr(torch.matrix_exp,
-           r"""
-Returns the matrix exponential. Supports batched input.
-For a matrix ``A``, the matrix exponential is defined as
+add_docstr(torch.matrix_exp, r"""
+matrix_exp(input) -> Tensor
+
+Computes the matrix exponential of a square matrix or of each square matrix in a batch.
+For a matrix :attr:`input`, the matrix exponential is defined as
 
 .. math::
-    \mathrm{e}^A = \sum_{k=0}^\infty A^k / k!
+    \mathrm{e}^\text{input} = \sum_{k=0}^\infty \text{input}^k / k!
 
 """ + r"""
 The implementation is based on: