pytorch
Write docstrings for PyTorch functions and methods following PyTorch conventions. Use when writing or updating docstrings in PyTorch code.
bunx add-skill pytorch/pytorch -s docstringLoading…
This skill describes how to write docstrings for functions and methods in the PyTorch project, following the conventions in torch/_tensor_docs.py and torch/nn/functional.py.
r"""...""") for all docstrings to avoid issues with LaTeX/math backslashesStart with the function signature showing all parameters:
r"""function_name(param1, param2, *, kwarg1=default1, kwarg2=default2) -> ReturnType
Notes:
* separator)Provide a one-line description of what the function does:
r"""conv2d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1) -> Tensor
Applies a 2D convolution over an input image composed of several input
planes.
Use Sphinx math directives for mathematical expressions:
.. math::
\text{Softmax}(x_{i}) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}
Or inline math: :math:\x^2``
Link to related classes and functions using Sphinx roles:
:class:\~torch.nn.ModuleName`` - Link to a class:func:\torch.function_name`` - Link to a function:meth:\~Tensor.method_name`` - Link to a method:attr:\attribute_name`` - Reference an attribute~ prefix shows only the last component (e.g., Conv2d instead of torch.nn.Conv2d)Example:
See :class:`~torch.nn.Conv2d` for details and output shape.
Use admonitions for important information:
.. note::
This function doesn't work directly with NLLLoss,
which expects the Log to be computed between the Softmax and itself.
Use log_softmax instead (it's faster and has better numerical properties).
.. warning::
:func:`new_tensor` always copies :attr:`data`. If you have a Tensor
``data`` and want to avoid a copy, use :func:`torch.Tensor.requires_grad_`
or :func:`torch.Tensor.detach`.
Document all parameters with type annotations and descriptions:
Args:
input (Tensor): input tensor of shape :math:`(\text{minibatch} , \text{in\_channels} , iH , iW)`
weight (Tensor): filters of shape :math:`(\text{out\_channels} , kH , kW)`
bias (Tensor, optional): optional bias tensor of shape :math:`(\text{out\_channels})`. Default: ``None``
stride (int or tuple): the stride of the convolving kernel. Can be a single number or a
tuple `(sH, sW)`. Default: 1
Formatting rules:
(Type), (Type, optional) for optional parametersvalue" at the end``None``Sometimes keyword arguments are documented separately:
Keyword args:
dtype (:class:`torch.dtype`, optional): the desired type of returned tensor.
Default: if None, same :class:`torch.dtype` as this tensor.
device (:class:`torch.device`, optional): the desired device of returned tensor.
Default: if None, same :class:`torch.device` as this tensor.
requires_grad (bool, optional): If autograd should record operations on the
returned tensor. Default: ``False``.
Document the return value:
Returns:
Tensor: Sampled tensor of same shape as `logits` from the Gumbel-Softmax distribution.
If ``hard=True``, the returned samples will be one-hot, otherwise they will
be probability distributions that sum to 1 across `dim`.
Or simply include it in the function signature line if obvious from context.
Always include examples when possible:
Examples::
>>> inputs = torch.randn(33, 16, 30)
>>> filters = torch.randn(20, 16, 5)
>>> F.conv1d(inputs, filters)
>>> # With square kernels and equal stride
>>> filters = torch.randn(8, 4, 3, 3)
>>> inputs = torch.randn(1, 4, 5, 5)
>>> F.conv2d(inputs, filters, padding=1)
Formatting rules:
Examples:: with double colon>>> prompt for Python code# when helpful>>>)Link to papers or external documentation:
.. _Link Name:
https://arxiv.org/abs/1611.00712
Reference them in text: See `Link Name`_
For regular Python functions, use a standard docstring:
def relu(input: Tensor, inplace: bool = False) -> Tensor:
r"""relu(input, inplace=False) -> Tensor
Applies the rectified linear unit function element-wise. See
:class:`~torch.nn.ReLU` for more details.
"""
# implementation
For C-bound functions, use _add_docstr:
conv1d = _add_docstr(
torch.conv1d,
r"""
conv1d(input, weight, bias=None, stride=1, padding=0, dilation=1, groups=1) -> Tensor
Applies a 1D convolution over an input signal composed of several input
planes.
See :class:`~torch.nn.Conv1d` for details and output shape.
Args:
input: input tensor of shape :math:`(\text{minibatch} , \text{in\_channels} , iW)`
weight: filters of shape :math:`(\text{out\_channels} , kW)`
...
""",
)
For in-place operations (ending with _), reference the original:
add_docstr_all(
"abs_",
r"""
abs_() -> Tensor
In-place version of :meth:`~Tensor.abs`
""",
)
For aliases, simply reference the original:
add_docstr_all(
"absolute",
r"""
absolute() -> Tensor
Alias for :func:`abs`
""",
)
Use LaTeX math notation for tensor shapes:
:math:`(\text{minibatch} , \text{in\_channels} , iH , iW)`
For commonly used arguments, define them once and reuse:
common_args = parse_kwargs(
"""
dtype (:class:`torch.dtype`, optional): the desired type of returned tensor.
Default: if None, same as this tensor.
"""
)
# Then use with .format():
r"""
...
Keyword args:
{dtype}
{device}
""".format(**common_args)
Insert reproducibility notes or other common text:
r"""
{tf32_note}
{cudnn_reproducibility_note}
""".format(**reproducibility_notes, **tf32_notes)
Here's a complete example showing all elements:
def gumbel_softmax(
logits: Tensor,
tau: float = 1,
hard: bool = False,
eps: float = 1e-10,
dim: int = -1,
) -> Tensor:
r"""
Sample from the Gumbel-Softmax distribution and optionally discretize.
Args:
logits (Tensor): `[..., num_features]` unnormalized log probabilities
tau (float): non-negative scalar temperature
hard (bool): if ``True``, the returned samples will be discretized as one-hot vectors,
but will be differentiated as if it is the soft sample in autograd. Default: ``False``
dim (int): A dimension along which softmax will be computed. Default: -1
Returns:
Tensor: Sampled tensor of same shape as `logits` from the Gumbel-Softmax distribution.
If ``hard=True``, the returned samples will be one-hot, otherwise they will
be probability distributions that sum to 1 across `dim`.
.. note::
This function is here for legacy reasons, may be removed from nn.Functional in the future.
Examples::
>>> logits = torch.randn(20, 32)
>>> # Sample soft categorical using reparametrization trick:
>>> F.gumbel_softmax(logits, tau=1, hard=False)
>>> # Sample hard categorical using "Straight-through" trick:
>>> F.gumbel_softmax(logits, tau=1, hard=True)
.. _Link 1:
https://arxiv.org/abs/1611.00712
"""
# implementation
When writing a PyTorch docstring, ensure:
r"""):func:, :class:, :meth:):class::class:\~torch.nn.Module`` - Class reference:func:\torch.function`` - Function reference:meth:\~Tensor.method`` - Method reference:attr:\attribute`` - Attribute reference:math:\equation`` - Inline math:ref:\label`` - Internal reference``code`` - Inline code (use double backticks)``True`` ``None`` ``False``Tensor, int, float, bool, str, tuple, list, etc.Use when you need to run Flow type checking, or when seeing Flow type errors in React code.
Use when you want to validate changes before committing, or when you need to check all React contribution requirements.
Use when feature flag tests fail, flags need updating, understanding @gate pragmas, debugging channel-specific test failures, or adding new flags to React.
Use when you need to check feature flag states, compare channels, or debug why a feature behaves differently across release channels.