Docstring

TorchRec Docstring Writing Guide

This skill describes how to write docstrings for functions and methods in the TorchRec project, following PyTorch conventions.

General Principles

• Use raw strings (r"""...""") for all docstrings to avoid issues with LaTeX/math backslashes
• Follow Sphinx/reStructuredText (reST) format for documentation
• Be concise but complete - include all essential information
• Always include examples when possible
• Use cross-references to related functions/classes

Docstring Structure

1. Function Signature (First Line)

Start with the function signature showing all parameters:

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:

r"""apply_optimizer_in_backward(optimizer_class, params, optimizer_kwargs) -> None

Applies optimizer to parameters in backward pass for memory efficiency.

3. Mathematical Formulas (if applicable)

Use Sphinx math directives for mathematical expressions:

.. math::
    \text{output} = \text{input} \cdot \text{weight}^T

Or inline math: :math:\x^2``

4. Cross-References

Link to related classes and functions using Sphinx roles:

• :class:\~torchrec.modules.EmbeddingBagCollection`` - Link to a class
• :func:\torchrec.distributed.sharding.shard`` - Link to a function
• :meth:\~Module.forward`` - Link to a method
• :attr:\attribute_name`` - Reference an attribute
• The ~ prefix shows only the last component

Example:

See :class:`~torchrec.distributed.DistributedModelParallel` for details.

5. Notes and Warnings

Use admonitions for important information:

.. note::
    This function requires CUDA to be available.

.. warning::
    This API is experimental and may change without notice.

6. Args Section

Document all parameters with type annotations and descriptions:

Args:
    module (nn.Module): Module to be sharded across devices.
    device (torch.device, optional): Device to place the module. Default: ``None``
    sharders (List[ModuleSharder], optional): List of sharders to use for sharding.
        Default: ``None``
    plan (ShardingPlan, optional): Explicit sharding plan. If not provided, will be
        generated automatically. Default: ``None``

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 4 spaces

7. Returns Section

Document the return value:

Returns:
    ShardedModule: The sharded module ready for distributed training.
        The module will have its parameters distributed according to
        the sharding plan.

8. Raises Section (if applicable)

Document exceptions that may be raised:

Raises:
    ValueError: If the sharding plan is invalid for the given module.
    RuntimeError: If CUDA is not available when GPU sharding is requested.

9. Examples Section

Always include examples when possible:

Examples::

    >>> import torchrec
    >>> from torchrec.modules import EmbeddingBagCollection
    >>> ebc = EmbeddingBagCollection(
    ...     tables=[
    ...         EmbeddingBagConfig(
    ...             name="product",
    ...             embedding_dim=64,
    ...             num_embeddings=1000,
    ...             feature_names=["product_id"],
    ...         ),
    ...     ],
    ... )
    >>> # Shard the module
    >>> sharded_ebc = shard(ebc, plan=plan)

Formatting rules:

• Use Examples:: with double colon
• Use >>> prompt for Python code
• Include comments with # when helpful
• Show actual output when it helps understanding

TorchRec-Specific Patterns

Embedding Configuration

Args:
    tables (List[EmbeddingBagConfig]): List of embedding table configurations.
        Each config specifies the table name, embedding dimension, number of
        embeddings, and feature names.
    device (Optional[torch.device]): Device to place embeddings. Default: ``None``

Sharding-Related

Args:
    sharding_type (ShardingType): How to shard the embedding table. Options are:
        - ``TABLE_WISE``: Each table on a single device
        - ``ROW_WISE``: Rows distributed across devices
        - ``COLUMN_WISE``: Columns distributed across devices
        - ``TABLE_ROW_WISE``: Combination of table and row sharding

KeyedJaggedTensor

Args:
    kjt (KeyedJaggedTensor): Sparse features in KeyedJaggedTensor format.
        Contains keys (feature names), values (embedding indices), and
        lengths/offsets for variable-length sequences.

Complete Example

def shard_modules(
    module: nn.Module,
    plan: ShardingPlan,
    env: ShardingEnv,
    device: Optional[torch.device] = None,
) -> nn.Module:
    r"""
    Shard a module's embedding tables according to a sharding plan.

    This function takes a module containing embedding tables and distributes
    them across multiple devices according to the provided sharding plan.
    It supports various sharding strategies including table-wise, row-wise,
    and column-wise sharding.

    Args:
        module (nn.Module): The module containing embedding tables to shard.
        plan (ShardingPlan): The sharding plan specifying how each table
            should be distributed.
        env (ShardingEnv): The sharding environment containing process group
            information and device topology.
        device (torch.device, optional): Target device for local shards.
            Default: ``None`` (uses current device)

    Returns:
        nn.Module: The sharded module with distributed embedding tables.

    Raises:
        ValueError: If the plan references tables not present in the module.
        RuntimeError: If the sharding environment is not properly initialized.

    .. note::
        This function modifies the module in-place for efficiency.

    .. warning::
        This is an experimental API and may change in future releases.

    Examples::

        >>> from torchrec.distributed import shard_modules
        >>> from torchrec.distributed.planner import EmbeddingShardingPlanner
        >>>
        >>> # Create a sharding plan
        >>> planner = EmbeddingShardingPlanner()
        >>> plan = planner.plan(module, sharders)
        >>>
        >>> # Shard the module
        >>> sharded_module = shard_modules(module, plan, env)
    """
    # implementation

Quick Checklist

When writing a TorchRec docstring, ensure:

• Use raw string (r""")
• Include function signature on first line
• Provide brief description
• Document all parameters in Args section with types
• Include default values for optional parameters
• Use Sphinx cross-references (:func:, :class:, :meth:)
• Add mathematical formulas if applicable
• Include at least one example in Examples section
• Add warnings/notes for experimental APIs
• Document any exceptions in Raises section
• Use proper math notation for tensor shapes