---
name: megatron-python-style
source: https://app.decimal.ai/s/megatron-python-style@1/SKILL.md
source_sha256: ea0edca678aa
---

# Megatron-LM Python style conventions

## Contract

Enforce the Megatron-LM Python style on every function, method, and class you emit:
modern `X | None` / builtin-generic annotations, Google-style docstrings, specific
exception catches, 119-char lines, and the standard naming casing. Apply to all Python
code written for this codebase — these are non-default choices the base model gets wrong.

## Rules

### 1. Optionals — `X | None`, never `Optional[X]`

- Write an optional value as `X | None`. NEVER `Optional[X]` and NEVER `Union[X, None]`.
- The `| None` must appear at the **end** of the annotation: `dict[str, int] | None`, not
  `None | dict[str, int]`.
- A parameter with a `None` default still gets the full annotation: `tags: set[str] | None = None`.

### 2. Unions — `A | B`, never `Union[A, B]`

- Write a union as `A | B | C`. NEVER `Union[A, B]`.
- Combine with optional as `A | B | None`.

### 3. Generics — builtin lowercase, never the `typing` aliases

- Use the builtin generics: `list[int]`, `dict[str, int]`, `tuple[int, ...]`, `set[str]`,
  `type[Foo]`, `frozenset[str]`.
- NEVER use `List`, `Dict`, `Tuple`, `Set`, `Type`, `FrozenSet` from `typing`.
- Nest them the same way: `list[list[int]]`, `dict[str, list[int]]`, `list[dict[str, int]]`.
- A homogeneous variable-length tuple is `tuple[int, ...]`; a fixed-shape tuple names each
  slot: `tuple[int, int]`, `tuple[str, float, bool]`.

### 4. Imports — the deprecated `typing` names must not appear

- Do NOT write `from typing import Optional, Union, List, Dict, Tuple, Set, Type`.
- Those symbols must not appear anywhere in the file — not in annotations, not in imports.
- `typing` is still fine for things with no builtin form (`Any`, `Callable`, `Protocol`,
  `TypeVar`, `Iterable`); only the six aliases above are banned.

### 5. Full type hints on every public callable

- Every public function and method carries a type hint on **every parameter** and on the
  **return**. A function that returns nothing is annotated `-> None`.
- `self` and `cls` are not annotated.

### 6. Docstrings — Google style on every public class/function/method

- One-line summary first, then a blank line, then the sections.
- `Args:` section listing each parameter by name.
- `Returns:` section describing the return value (omit only when the function returns `None`
  and that is obvious; prefer including it).
- `Raises:` section listing each exception type the body can raise — present **if and only
  if** the body can raise.
- Classes get a class-level docstring describing the class; their public methods get their own.

### 7. Exceptions — catch specific types, never a bare or catch-all except

- NEVER write a bare `except:`.
- NEVER write `except Exception:` as the normal catch-all.
- Catch the specific type(s): `except KeyError:`, `except (ValueError, TypeError):`,
  `except FileNotFoundError:`.

### 8. Formatting & naming

- Line length is **119** characters — not 79, not 88, not 100. Wrap only past 119.
- `snake_case` for functions and variables, `PascalCase` for classes, `UPPER_SNAKE` for
  module-level constants.

## Worked examples

### Rule 1 — optionals

BEFORE (base default):
```python
from typing import Optional
def find_checkpoint(run_dir: str, step: int) -> Optional[str]:
    ...
```
AFTER (conforming):
```python
def find_checkpoint(run_dir: str, step: int) -> str | None:
    ...
```

### Rule 2 — unions

BEFORE:
```python
from typing import Union
def coerce(x: Union[int, float]) -> Union[int, None]:
    ...
```
AFTER:
```python
def coerce(x: int | float) -> int | None:
    ...
```

### Rule 3 — builtin generics

BEFORE:
```python
from typing import List, Dict
def merge_shards(shards: List[int]) -> Dict[str, object]:
    ...
```
AFTER:
```python
def merge_shards(shards: list[int]) -> dict[str, object]:
    ...
```

### Rule 4 — no banned imports

BEFORE:
```python
from typing import Optional, List, Dict, Tuple
```
AFTER (the line is simply gone; annotations use `| None`, `list`, `dict`, `tuple`):
```python
# no typing import of the aliased names at all
```

### Rule 5 + 6 — full hints and Google docstring

BEFORE:
```python
def batch_tokenize(texts, max_len):
    """Tokenize texts."""
    ...
```
AFTER:
```python
def batch_tokenize(texts: list[str], max_len: int) -> list[list[int]]:
    """Tokenize a batch of strings into token-id lists.

    Args:
        texts: Input strings to tokenize.
        max_len: Maximum number of tokens per sequence.

    Returns:
        One token-id list per input string.
    """
    ...
```

### Rule 6 — `Raises:` when the body can raise

BEFORE:
```python
def validate_shape(shape: tuple[int, ...]) -> tuple[int, ...]:
    """Validate a shape."""
    if any(d <= 0 for d in shape):
        raise ValueError("dimensions must be positive")
    return shape
```
AFTER:
```python
def validate_shape(shape: tuple[int, ...]) -> tuple[int, ...]:
    """Validate that every dimension of a shape is positive.

    Args:
        shape: The tensor shape to validate.

    Returns:
        The shape unchanged when valid.

    Raises:
        ValueError: If any dimension is non-positive.
    """
    if any(d <= 0 for d in shape):
        raise ValueError("dimensions must be positive")
    return shape
```

### Rule 7 — specific exception catch

BEFORE:
```python
def parse_config(path: str):
    try:
        return json.load(open(path))
    except:
        return None
```
AFTER:
```python
def parse_config(path: str) -> dict | None:
    """Read and parse a JSON config file.

    Args:
        path: Path to the config file.

    Returns:
        The parsed config, or None if the file is missing.
    """
    try:
        with open(path) as fh:
            return json.load(fh)
    except FileNotFoundError:
        return None
```

## Edge cases & exceptions

- **`None` default, optional type:** annotate as `set[str] | None = None`. The default
  value `None` does not let you drop the annotation.
- **Fixed vs. variable tuple:** `tuple[int, int]` (exactly two ints) is different from
  `tuple[int, ...]` (any number). Pick the one the function actually produces.
- **Forward refs / self-types:** still use `| None` inside the quoted form,
  e.g. `def next(self) -> "Node | None":`.
- **`typing.Any`, `Callable`, `Protocol`, `TypeVar`** are still imported from `typing` —
  only the six collection/optional aliases are banned.
- **Re-raising:** `except SpecificError: raise` is fine; the ban is on catching broadly,
  not on re-raising a specific type.
- **Truly want every error?** Even then, prefer the narrowest superclass that fits
  (e.g. `except OSError:`) over `except Exception:`. A bare `except:` is never allowed.
- **Private one-liner helpers** still take the modern annotations; the docstring rule is
  strongest for public API but the typing/exception rules apply everywhere.

## Do / Don't

- DON'T `from typing import Optional`; ALWAYS write `X | None` inline.
- DON'T annotate `Dict[str, int]`; ALWAYS write `dict[str, int]`.
- DON'T write `except:` or `except Exception:`; ALWAYS name the exact exception type.
- DON'T leave a public function with an untyped parameter; ALWAYS hint every parameter and
  the return.
- DON'T write a `Raises:` section for a function that cannot raise; ALWAYS add one when it can.
- DON'T wrap lines at 79/88/100; ALWAYS allow up to 119.
- DON'T use `camelCase` methods or `lowercase` class names; ALWAYS `snake_case` /
  `PascalCase` / `UPPER_SNAKE`.

## Common mistakes

- Reaching for `Optional[...]` / `Union[...]` out of habit — the single most common miss.
- Capitalized generics `List[int]` / `Dict[...]` with a `from typing import ...` line.
- A bare `except:` or `except Exception: pass` to "be safe."
- A docstring that is just a one-line summary with no `Args:` / `Returns:`.
- Omitting the return annotation (especially `-> None`).
- Adding a `Raises:` section out of habit when the function never raises, or omitting it
  when the function clearly does.

## Quick checklist

1. No `Optional[` / `Union[` anywhere — `X | None`, `A | B`.
2. `list/dict/tuple/set/type[...]`, never the `typing` aliases; no banned imports.
3. Every public param and return annotated; `-> None` when nothing is returned.
4. Google docstring: summary, blank line, `Args:`, `Returns:`, `Raises:` (iff it raises).
5. Specific `except`; never bare or `except Exception:`.
6. ≤119-char lines; `snake_case` / `PascalCase` / `UPPER_SNAKE`.