Docs fix

[pggPL]

Description

Our documentation returned a lot of warnings and it seems that some of them were rational. It turned out that half of our PyTorch API was not rendered. This PR fixes all the warnings and forces Github workflow to error out if any docs warning will appear.

The most problematic error was related to cyclic imports - this resulted in part of our PyTorch API not being rendered. Other were mostly related to wrong formatting and some warnings didn't result in anything wrong.

Our docs for 2.9 contains only part of PyTorch API exposed in 2.8, so this will need urgent fix I will update in separate PR.

Type of change

• Documentation change (change only to the documentation, either a fix or a new content)
• Bug fix (non-breaking change which fixes an issue)
• New feature (non-breaking change which adds functionality)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• Infra/Build change
• Code refactoring

Checklist:

• [ x] I have read and followed the contributing guidelines
• The functionality is complete
• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes

[greptile-apps]

Greptile Overview

Greptile Summary

This PR resolves critical documentation rendering issues where approximately half of the PyTorch API was not being rendered in the generated documentation. The root causes were cyclic import dependencies and reStructuredText formatting violations. The fix involves restructuring imports across JAX and PyTorch modules to break circular dependencies (primarily moving QuantizeLayout and torch_version imports to their source modules), correcting RST section underlines to match header text lengths, converting docstring formatting to proper RST syntax, and enforcing strict documentation builds in CI by adding the -W flag to treat warnings as errors.

Important Files Changed

Filename	Score	Overview
.github/workflows/docs.yml	5/5	Added -W flag to Sphinx builds to make documentation warnings fatal in CI
transformer_engine/pytorch/utils.py	5/5	Converted torch_version from module-level import to cached function to break cyclic dependencies
transformer_engine/pytorch/cross_entropy.py	5/5	Converted parallel_cross_entropy from bare function reference to documented wrapper function
transformer_engine/jax/quantize/quantizer.py	5/5	Removed QuantizeLayout from __all__ to break circular dependency chain
transformer_engine/jax/quantize/hadamard.py	5/5	Changed QuantizeLayout import from delayed local import to top-level C++ extension import
transformer_engine/jax/cpp_extensions/misc.py	5/5	Moved QuantizeLayout import from relative path to C++ extension to break cycle
transformer_engine/jax/cpp_extensions/activation.py	5/5	Moved QuantizeLayout import from quantize module to C++ extension module
transformer_engine/jax/cpp_extensions/quantization.py	5/5	Moved QuantizeLayout import from quantize module to C++ extension module
transformer_engine/jax/cpp_extensions/gemm.py	5/5	Moved QuantizeLayout import from quantize module to C++ extension module
transformer_engine/jax/cpp_extensions/normalization.py	5/5	Moved QuantizeLayout import from quantize module to C++ extension module
transformer_engine/jax/dense.py	5/5	Moved QuantizeLayout import from local module to top-level package
transformer_engine/pytorch/jit.py	5/5	Changed torch_version import from current module to utils submodule
transformer_engine/pytorch/distributed.py	5/5	Changed torch_version import from package-level to explicit utils module
transformer_engine/pytorch/quantization.py	5/5	Reorganized imports with proper blank line separation for clarity
transformer_engine/pytorch/module/linear.py	5/5	Changed torch_version import to come from utils submodule
transformer_engine/pytorch/module/layernorm_linear.py	5/5	Changed torch_version import to come from utils submodule
transformer_engine/pytorch/module/layernorm_mlp.py	4/5	Changed torch_version import and improved docstring RST formatting
transformer_engine/pytorch/ops/_common.py	5/5	Changed torch_version import to come from utils submodule
transformer_engine/pytorch/ops/basic/l2normalization.py	5/5	Changed torch_version import to come from utils submodule
transformer_engine/pytorch/transformer.py	5/5	Fixed torch_version import and improved docstring RST formatting
transformer_engine/pytorch/module/base.py	5/5	Converted docstring code blocks from markdown to RST syntax
transformer_engine/pytorch/attention/dot_product_attention/dot_product_attention.py	5/5	Improved RST formatting of docstring with proper inline code and bullet lists
transformer_engine/pytorch/attention/multi_head_attention.py	5/5	Improved RST formatting of docstring with proper inline code and bullet lists
transformer_engine/jax/flax/transformer.py	5/5	Converted ASCII table in docstring to RST table directive
docs/conf.py	4/5	Added custom logging filter to suppress unavoidable duplicate namespace warnings
docs/api/pytorch.rst	4/5	Reorganized API documentation with new sections and deprecated functions section
docs/debug.rst	5/5	Added blank line after copyright header to fix RST parsing
docs/debug/api.rst	5/5	Added blank line after copyright header to fix RST parsing
docs/debug/3_api_features.rst	5/5	Extended section underline to match header length
docs/debug/3_api_debug_setup.rst	5/5	Extended section underlines to match header lengths
docs/debug/4_distributed.rst	5/5	Extended section underlines to match header lengths
docs/debug/2_config_file_structure.rst	5/5	Extended section underlines to match header lengths
docs/debug/1_getting_started.rst	0/5	Section underlines still don't match header text lengths exactly - will cause RST warnings
docs/examples/attention/attention.ipynb	5/5	Updated API reference links to lowercase format for Sphinx compatibility
docs/examples/te_gemma/tutorial_generation_gemma_with_te.ipynb	5/5	Fixed broken internal documentation link

Confidence score: 2/5

• This PR contains one file (docs/debug/1_getting_started.rst) with incorrect RST section underline lengths that will cause Sphinx warnings, directly contradicting the PR's goal of eliminating all warnings
• All other changes are formatting and import refactoring with minimal risk, but the incomplete fix in 1_getting_started.rst will cause the new CI check (which treats warnings as errors) to fail
• Pay close attention to docs/debug/1_getting_started.rst - the underlines need to match the exact character length of each section header, not use a fixed-width approach

_{35 files reviewed, 8 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

syntax: underline mismatch - title has 15 chars ('Getting started') but underline has 31 '=' chars

Prompt To Fix With AI

This is a comment left during a code review.
Path: docs/debug/1_getting_started.rst
Line: 7:7

Comment:
**syntax:** underline mismatch - title has 15 chars ('Getting started') but underline has 31 '=' chars

How can I resolve this? If you propose a fix, please make it concise.

[greptile-apps]

syntax: underline mismatch - title has 24 chars ('Example training script') but underline has 30 '-' chars

Prompt To Fix With AI

This is a comment left during a code review.
Path: docs/debug/1_getting_started.rst
Line: 41:41

Comment:
**syntax:** underline mismatch - title has 24 chars ('Example training script') but underline has 30 '-' chars

How can I resolve this? If you propose a fix, please make it concise.

[greptile-apps]

syntax: underline mismatch - title has 11 chars ('Config file') but underline has 30 '-' chars

Prompt To Fix With AI

This is a comment left during a code review.
Path: docs/debug/1_getting_started.rst
Line: 84:84

Comment:
**syntax:** underline mismatch - title has 11 chars ('Config file') but underline has 30 '-' chars

How can I resolve this? If you propose a fix, please make it concise.

[greptile-apps]

syntax: underline mismatch - title has 22 chars ('Adjusting Python file') but underline has 28 '-' chars

Prompt To Fix With AI

This is a comment left during a code review.
Path: docs/debug/1_getting_started.rst
Line: 117:117

Comment:
**syntax:** underline mismatch - title has 22 chars ('Adjusting Python file') but underline has 28 '-' chars

How can I resolve this? If you propose a fix, please make it concise.

[greptile-apps]

syntax: underline mismatch - title has 20 chars ('Inspecting the logs') but underline has 28 '-' chars

Prompt To Fix With AI

This is a comment left during a code review.
Path: docs/debug/1_getting_started.rst
Line: 149:149

Comment:
**syntax:** underline mismatch - title has 20 chars ('Inspecting the logs') but underline has 28 '-' chars

How can I resolve this? If you propose a fix, please make it concise.

[greptile-apps]

syntax: underline mismatch - title has 27 chars ('Logging using TensorBoard') but underline has 28 '-' chars

Prompt To Fix With AI

This is a comment left during a code review.
Path: docs/debug/1_getting_started.rst
Line: 218:218

Comment:
**syntax:** underline mismatch - title has 27 chars ('Logging using TensorBoard') but underline has 28 '-' chars

How can I resolve this? If you propose a fix, please make it concise.

[greptile-apps]

syntax: trailing whitespace on line 109

Prompt To Fix With AI

This is a comment left during a code review.
Path: docs/conf.py
Line: 109:110

Comment:
**syntax:** trailing whitespace on line 109

How can I resolve this? If you propose a fix, please make it concise.

[greptile-apps]

Greptile Overview

Greptile Summary

This review covers only the most recent changes that address syntax and formatting issues identified in previous reviews. The developer has made targeted fixes to resolve trailing whitespace, RST heading underline mismatches, and missing newlines across several documentation files. These changes are part of the broader documentation quality enforcement effort.

The recent commits address:

• Trailing whitespace in docs/conf.py (line 109)
• RST heading underline length mismatches in docs/debug/1_getting_started.rst (6 sections)
• Missing newline at end of file in docs/api/pytorch.rst

Important Files Changed

Filename	Score	Overview
docs/conf.py	5/5	Fixed trailing whitespace on line 109
docs/debug/1_getting_started.rst	3/5	Adjusted RST heading underlines but mismatches remain between title lengths and underline lengths
docs/api/pytorch.rst	5/5	Added missing newline at end of file

Confidence score: 3/5

• This PR contains syntax fixes but incomplete RST formatting corrections that will still cause Sphinx warnings
• Score reflects that while trailing whitespace and missing newlines are fixed, the RST heading underlines in docs/debug/1_getting_started.rst still don't match their title character counts (e.g., "Getting started" has 15 chars but 31 '=' underline chars), which will cause Sphinx build failures when the -W flag is active
• Pay close attention to docs/debug/1_getting_started.rst - all six section heading underlines need to be adjusted to exactly match their title character counts before the CI will pass

Additional Comments (1)

1. docs/debug/3_api_debug_setup.rst, line 89 (link)

   syntax: missing newline at end of file

_{34 files reviewed, 12 comments}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

style: : blank line between imports - typically import groups (stdlib, third-party, local) are separated by blank lines, but mixing local imports like this breaks the pattern

Suggested change

	from .constants import dist_group_type
	from .constants import dist_group_type
	from .utils import get_device_compute_capability

[greptile-apps]

Greptile Overview

Greptile Summary

This review covers only the changes made since the last review, not the entire PR. The developer has successfully resolved all remaining ReStructuredText (RST) formatting issues across five debug documentation files. These changes correct section underline mismatches where underline lengths did not match their corresponding title lengths—a strict RST syntax requirement that causes Sphinx documentation build warnings. The fixes affect files in docs/debug/ including 1_getting_started.rst, 3_api_debug_setup.rst, 3_api_features.rst, api.rst, and 4_distributed.rst. Each correction ensures that section titles and their underlines have exactly matching character counts (e.g., "API" with 3 '=' chars, "Getting started" with 15 '=' chars), which eliminates documentation warnings and ensures proper rendering of the PyTorch API documentation. These changes are part of a broader effort to fix all documentation warnings and enforce clean doc builds in CI.

Important Files Changed

Filename	Score	Overview
docs/debug/3_api_debug_setup.rst	5/5	Fixed RST underline lengths for three function headers (initialize(), set_tensor_reduction_group(), set_weight_tensor_tp_group_reduce())
docs/debug/3_api_features.rst	5/5	Corrected section underline from 27 to 14 '=' chars to match "Debug features" title
docs/debug/api.rst	5/5	Fixed "API" section underline from 12 to 3 '=' chars
docs/debug/4_distributed.rst	5/5	Corrected underlines for main title and four subsections throughout the distributed training documentation
docs/debug/1_getting_started.rst	5/5	Fixed six section underlines to match title lengths across getting started guide

Confidence score: 5/5

• This PR is safe to merge with minimal risk—all changes are pure RST formatting corrections with no functional impact
• Score reflects that these are syntax-only corrections (underline length matching) with zero functional changes, no logic modifications, and no code execution paths affected
• No files require special attention—all changes follow the same trivial pattern of correcting underline character counts to match section titles

_{5 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}

[ptrendx]

We do want to have them still in the documentation though, just marked as deprecated.

[ptrendx]

Ok, I see that you just moved it.

[ptrendx]

In the past I had some trouble with the titles where the underline was not exact same length as the text.

[pggPL]

If the line is shorter, then warning is generated. In that case there is no warning and it renders properly, so I guess it's ok now.

[ptrendx]

Are you sure about that? Just looked at the current docs page and the link works properly there.

[pggPL]

Oh, some debugging leftover. It does also work, but I will remove this change.

[ptrendx]

Why _input and not just input?

[pggPL]

I also noticed this. This function already uses _input as an argument name, so I documented what is in the code rather than changing that - it would be breaking change. This function does not render now due to error btw.

[greptile-apps]

Greptile Overview

Greptile Summary

This PR fixes critical documentation build warnings that prevented half of the PyTorch API from being rendered in v2.9 documentation. The changes resolve cyclic import issues, improve RST formatting, and enforce warning-free builds in CI.

Key changes:

• Fixed cyclic import by moving torch_version from a cached function in __init__.py to a module-level variable in utils.py, with updated imports across affected modules
• Reorganized docs/api/pytorch.rst with logical sections (Recipe availability, MoE functions, Communication-computation overlap, Deprecated functions) for better API discoverability
• Improved RST docstring formatting with proper code literals (backticks), bullet lists, and table syntax for better Sphinx rendering
• Added SPHINXOPTS="-W" to GitHub workflow to fail builds on warnings, preventing future documentation regressions
• Added custom logging filter in conf.py to suppress unavoidable duplicate C++ namespace warnings
• Cleared Jupyter notebook execution counts to reduce version control noise

Confidence Score: 5/5

• This PR is safe to merge with minimal risk - it fixes documentation issues without affecting runtime behavior
• All changes are focused on documentation quality improvements and formatting fixes. The cyclic import fix is well-executed with consistent updates across all dependent modules. The workflow enforcement ensures future documentation quality. Only minor formatting issue found (missing newline)
• No files require special attention - only a trivial missing newline in docs/api/pytorch.rst

Important Files Changed

File Analysis

Filename	Score	Overview
.github/workflows/docs.yml	5/5	Added -W flag to Sphinx build to error on warnings, enforcing documentation quality
docs/conf.py	5/5	Added custom logging filter to suppress duplicate C++ namespace warnings, updated exclude patterns
transformer_engine/pytorch/init.py	5/5	Moved torch_version from function to module-level variable, fixing cyclic import issues
docs/api/pytorch.rst	4/5	Reorganized API docs with sections, moved deprecated functions, fixed rendering issues

Sequence Diagram

sequenceDiagram
    participant Dev as Developer
    participant GHA as GitHub Actions
    participant Sphinx as Sphinx Builder
    participant Code as Python Code
    participant Docs as Documentation

    Dev->>Code: Fix cyclic imports (torch_version)
    Dev->>Code: Update import paths
    Dev->>Docs: Reorganize API structure
    Dev->>Docs: Fix RST formatting issues
    Dev->>GHA: Add SPHINXOPTS="-W" flag
    
    Note over Dev,Docs: PR Changes Applied
    
    GHA->>Sphinx: Trigger docs build with -W
    Sphinx->>Code: Import Python modules
    Code-->>Sphinx: Successfully imported (no cycles)
    Sphinx->>Docs: Parse RST files
    Docs-->>Sphinx: Valid RST syntax
    Sphinx->>Sphinx: Apply custom warning filter
    Sphinx-->>GHA: Build succeeds (no warnings)
    
    Note over GHA: Future builds will error<br/>if warnings appear

Loading

_{9 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}

[greptile-apps]

syntax: missing final newline

Suggested change

	.. autoapifunction:: transformer_engine.pytorch.fp8_model_init
	.. autoapifunction:: transformer_engine.pytorch.fp8_model_init

[greptile-apps]

Greptile Overview

Greptile Summary

Fixed critical documentation build issues that prevented half of the PyTorch API from rendering in version 2.9. Resolved cyclic import in pytorch/__init__.py by moving torch_version() to utils.py, and enforced documentation quality by adding -W flag to Sphinx build.

• Cyclic import resolution: Moved torch_version() from pytorch/__init__.py to pytorch/utils.py to break import cycle that prevented API documentation generation
• Sphinx enforcement: Added -W flag to GitHub workflow to treat documentation warnings as build errors
• RST formatting fixes: Corrected underline mismatches in debug documentation files (15+ instances across multiple files)
• Docstring improvements: Enhanced formatting with proper RST syntax (backticks, bullet lists) in attention, cross_entropy, and other modules
• API documentation reorganization: Restructured pytorch.rst with logical sections and moved deprecated functions to the end
• Import corrections: Fixed several import issues in JAX and PyTorch modules (e.g., QuantizeLayout in JAX, torch_version references)
• Warning suppression: Added custom logging filter for unavoidable duplicate C++ declaration warnings

Confidence Score: 5/5

• This PR is safe to merge with minimal risk - it fixes critical documentation issues without changing runtime behavior
• All changes are documentation-focused (RST formatting, docstrings, import reorganization). The cyclic import fix is a proper refactoring that maintains the same API surface. The workflow change enforces quality going forward. No logic changes to core functionality.
• No files require special attention

Important Files Changed

File Analysis

Filename	Score	Overview
.github/workflows/docs.yml	5/5	Added -W flag to make Sphinx treat warnings as errors, enforcing documentation quality
docs/conf.py	5/5	Added custom logging filter to suppress duplicate C++ declaration warnings for transformer_engine namespace
transformer_engine/pytorch/init.py	5/5	Fixed cyclic import by removing torch_version() function from __init__.py (moved to utils.py)
transformer_engine/pytorch/utils.py	5/5	Added torch_version() function here to resolve cyclic import issues
docs/api/pytorch.rst	5/5	Reorganized API documentation with sections, moved deprecated functions to the end, added missing newline
transformer_engine/pytorch/cross_entropy.py	5/5	Added explicit function wrapper with full docstring for parallel_cross_entropy to improve API documentation

Sequence Diagram

sequenceDiagram
    participant Dev as Developer
    participant GH as GitHub Workflow
    participant Sphinx as Sphinx Builder
    participant Init as pytorch/__init__.py
    participant Utils as pytorch/utils.py
    
    Dev->>GH: Push docs changes
    GH->>Sphinx: Run build with -W flag
    
    Note over Sphinx: Before: Warnings ignored
    Note over Sphinx: After: Warnings = Errors
    
    Sphinx->>Sphinx: Check RST formatting
    Sphinx->>Sphinx: Check docstring syntax
    
    Note over Init,Utils: Cyclic Import Fix
    Init->>Utils: torch_version() moved here
    Utils->>Utils: Function callable via import
    
    Sphinx->>Sphinx: Parse API docs successfully
    Sphinx->>GH: Build succeeds (no warnings)
    GH->>Dev: CI passes

Loading

_{35 files reviewed, no comments}

_{Edit Code Review Agent Settings | Greptile}