Write docstrings for PyTorch functions and methods following PyTorch conventions. Use when writing or updating docstrings in PyTorch code.
npx skills add https://github.com/pytorch/pytorch --skill docstring
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.Statistical models library for Python. Use when you need specific model classes (OLS, GLM, mixed models, ARIMA) with detailed diagnostics, residuals, and inference. Best for econometrics, time series, rigorous inference with coefficient tables. For guided statistical test selection with APA reporting use statistical-analysis.
Answer questions about the AI SDK and help build AI-powered features. Use when developers: (1) Ask about AI SDK functions like generateText, streamText, ToolLoopAgent, embed, or tools, (2) Want to build AI agents, chatbots, RAG systems, or text generation features, (3) Have questions about AI providers (OpenAI, Anthropic, Google, etc.), streaming, tool calling, structured output, or embeddings, (4) Use React hooks like useChat or useCompletion. Triggers on: "AI SDK", "Vercel AI SDK", "generateText", "streamText", "add AI to my app", "build an agent", "tool calling", "structured output", "useChat".
Create an llms.txt file from scratch based on repository structure following the llms.txt specification at https://llmstxt.org/
Use when working directly with the `esm` Python SDK, ESM3 or ESMC model IDs, Forge/Biohub inference clients, or ESMFold2 folding workflows.
Modal is a serverless cloud platform for running Python on demand, including on-demand GPUs. Use when deploying or serving AI/ML models, running GPU-accelerated workloads (training, fine-tuning, inference), serving web endpoints, scheduling batch jobs, or scaling Python code to cloud containers with the Modal SDK.
Use Therapeutics Data Commons through the PyTDC Python package for registry discovery, approved dataset access, task-aware splits, evaluator metrics, benchmark groups, and bounded molecular-oracle workflows.
Machine learning in Python with scikit-learn. Use when working with supervised learning (classification, regression), unsupervised learning (clustering, dimensionality reduction), model evaluation, hyperparameter tuning, preprocessing, or building ML pipelines. Provides comprehensive reference documentation for algorithms, preprocessing techniques, pipelines, and best practices.
Statistical models library for Python. Use when you need specific model classes (OLS, GLM, mixed models, ARIMA) with detailed diagnostics, residuals, and inference. Best for econometrics, time series, rigorous inference with coefficient tables. For guided statistical test selection with APA reporting use statistical-analysis.
Take pytorch/docstring from the repository into ~/.claude/skills for personal
use, or into .claude/skills inside a project.
The agent identifies a skill by the name field in its header. Two skills with the
same name cannot sit side by side — one of them will be ignored.