1. Function Signature (First Line)
Start with the function signature showing all parameters:
```python
r"""function_name(param1, param2, *, kwarg1=default1, kwarg2=default2) -> ReturnType
```
Notes:
- Include the function name
- Show positional and keyword-only arguments (use
* separator) - Include default values
- Show return type annotation
- This line should NOT end with a period
2. Brief Description
Provide a one-line description of what the function does:
```python
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.
```
3. Mathematical Formulas (if applicable)
Use Sphinx math directives for mathematical expressions:
```python
.. math::
\text{Softmax}(x_{i}) = \frac{\exp(x_i)}{\sum_j \exp(x_j)}
```
Or inline math: :math:\x^2\``
4. Cross-References
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- The
~ prefix shows only the last component (e.g., Conv2d instead of torch.nn.Conv2d)
Example:
```python
See :class:~torch.nn.Conv2d for details and output shape.
```
5. Notes and Warnings
Use admonitions for important information:
```python
.. 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.
```
6. Args Section
Document all parameters with type annotations and descriptions:
```python
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:
- Parameter name in lowercase
- Type in parentheses:
(Type), (Type, optional) for optional parameters - Description follows the type
- For optional parameters, include "Default: `
value`" at the end - Use double backticks for inline code: ``
None `` - Indent continuation lines by 2 spaces
7. Keyword Args Section (if applicable)
Sometimes keyword arguments are documented separately:
```python
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`.
```
8. Returns Section (if needed)
Document the return value:
```python
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.
9. Examples Section
Always include examples when possible:
```python
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:
- Use
Examples:: with double colon - Use
>>> prompt for Python code - Include comments with
# when helpful - Show actual output when it helps understanding (indent without
>>>)
10. External References
Link to papers or external documentation:
```python
.. _Link Name:
https://arxiv.org/abs/1611.00712
```
Reference them in text: ``See Link Name_``