MegaBlocks

Public APIs (those without a leading underscore in the path) must follow these rules:

1. Documentation: All public APIs must have a docstring.
2. Type Annotations: All parameters must have type annotations.
3. Minimize Imports: Use native Python or PyTorch types for parameters whenever possible. It is acceptable to use a Union where one option is a primitive to allow users to pass strings instead of custom objects.

   # Example: allowing a string or a custom Device object
   def __init__(self, device: Union[str, Device]):
       if isinstance(device, str):
           device = Device(device)

4. Sequence Parameters: If a parameter accepts a sequence, allow None or a singleton to simplify the API. Use helpers like ensure_tuple to normalize inputs.

   from typing import Optional, Sequence, Union
   from torch import Tensor
   from composer.utils import ensure_tuple
   
   def foo(x: Optional[Union[Tensor, Sequence[Tensor]]]) -> tuple[Tensor, ...]:
       return ensure_tuple(x)

For non-core dependencies (e.g., specific models or algorithms), follow these steps to keep the base installation minimal:

1. Specify in setup.py: Add the dependency to the extra_deps dictionary. Users can then install it via pip install 'megablocks[extra_name]'.
2. Conditional Imports: Use a try/except block with MissingConditionalImportError to provide clear error messages if the dependency is missing.

from composer import Callback
from composer.utils import MissingConditionalImportError

class SystemMetricsMonitor(Callback):
    try:
        import pynvml
    except ImportError as e:
        raise MissingConditionalImportError(
            extra_deps_group="system_metrics_monitor",
            conda_package="pynvml",
            conda_channel="conda-forge",
        ) from e

If the dependency is core to MegaBlocks, add it to install_requires in setup.py and requirements.run in meta.yaml.

MegaBlocks uses .. doctest or .. testcode directives to provide executable examples in docstrings. These are verified in CI/CD to ensure documentation remains accurate.

Writing Doctests

• Use .. testcode:: for examples shown in the docs. Do not use .. code-block:: for Python examples as they are not tested.
• Use .. testsetup:: for setup code that should run before the test but not be displayed in the documentation.
• For global test fixtures, use docs/source/doctest_fixtures.py. For specific setup, use .. testsetup:: to avoid polluting the global namespace.

Running Doctests

You must complete an HTML build before running doctests to ensure all tests are identified.

1. Activate your virtual environment.
2. Navigate to the docs folder.
3. Run make clean.
4. Run make html.
5. Run make doctest.

import torch
from typing import Optional

def my_function(x: Optional[torch.Tensor]) -> torch.Tensor:
    """blah function

    Args:
        input (torch.Tensor): Your guess.

    Returns:
        torch.Tensor: How good your input is.

    Raises:
        ValueError: If your input is negative.

    Example:
        .. testsetup::

            # optional setup section, not shown in docs
            import torch
            x = torch.randn(42)


        .. testcode::

            # shown in docs; runs after testsetup
            my_function(x)
    """
    ...

source path/to/megablocks_venv/bin/activate
cd megablocks/docs
make clean
make html
make doctest 2>/dev/null # For more verbosity, do not direct stderr to /dev/null

MegaBlocks uses Google Style Docstrings. All public APIs must be documented.

Key formatting rules:

1. Content: Include a summary of the function/class, arguments, return statements (if not None), and custom exceptions.
2. Class __init__: Document arguments for the __init__ signature under the class-level docstring. Do not create a separate docstring for __init__.
3. Argument Annotations: Include the type. If an argument has a default value, specify optional in the type annotation and state the default value in the description.
4. Returns: Document the return type and description. For multiple return values (e.g., tuples), document each element individually.

from typing import Optional, Union

def foo(bar: int):
    """Foo.

    Args:
        bar (int): Required bar.
    """
    ...

def foo2(bar: int = 42):
    """Foo2.

    Args:
        bar (int, optional): The first Argument. Default: ``42``.
    """
    ...

def foo3(bar: Optional[int] = None):
    """Foo3.

    Args:
        bar (int, optional): The first Argument. Default: ``None``.
    """
    ...

def foo4(bar: Union[int, str] = 42):
    """Foo4.

    Args:
        bar (int | str, optional): The first Argument. Default: ``42``.
    """
    ...

def foo5(bar: int) -> int:
    """Foo5.

    Args:
        bar (int): Required bar.

    Returns:
        int: Description of return statement.
    """
    ...

def foo6(bar: int) -> tuple[int, str]:
    """Foo6.

    Args:
        bar (int): Required bar.

    Returns:
        a (int): Returned value.
        b (str): Returned value.
    """
    ...