[Docs] Add automatic operation documentation generation with PyTorch-style structure

[samiyadev786]

Summary

This PR implements automatic operation discovery and documentation generation from forge/forge/op/*.py files, addressing the requirement for automatic discovery of operations similar to Doxygen/PyTorch documentation style.

Changes

🎯 Core Implementation

1. Automatic Operation Discovery (scripts/discover_operations.py)

   • New module that automatically scans forge/forge/op/*.py files
   • Uses Python AST to parse function definitions
   • Extracts function signatures, docstrings, and parameter information
   • Infers operation categories from file names
   • Discovers 78 operations automatically from the codebase

2. Enhanced Documentation Generator (scripts/generate_ops_docs.py)

   • Integrated automatic discovery as the primary method
   • Uses actual Python function signatures from source code
   • Improved docstring parsing with better edge case handling
   • Enhanced parameter description formatting (removes redundant type prefixes)
   • Fixed incorrect descriptions (e.g., "Sigmoid" for Abs → "Elementwise absolute value operation")
   • Cleaner return value descriptions
   • Falls back to ops_data.py only if automatic discovery fails

📝 Generated Documentation

3. Operation Documentation Pages

   • Generated 78 operation pages using forge.op.* naming convention
   • All pages follow PyTorch-style documentation format
   • Consistent structure with accurate function signatures
   • Each page includes function signature, parameters, returns, and descriptions

4. Index Page (docs/src/operations.md)

   • Categorized list of all 78 operations
   • Operations grouped by type (Elementwise, Convolution, Pooling, etc.)
   • Short descriptions for each operation
   • Links to individual operation pages

Features

✅ Automatic Discovery: Operations are automatically discovered from Python source files
✅ No Manual Entry Required: Simply define operations with proper docstrings in forge/forge/op/*.py
✅ Accurate Signatures: Uses actual Python function signatures from source code
✅ Better Formatting: Cleaner parameter and return descriptions
✅ Category Inference: Automatically categorizes operations based on file location
✅ Edge Case Handling: Handles missing/incorrect docstrings gracefully
✅ PyTorch-Style Structure: Matches PyTorch documentation format with root page and individual pages

How It Works

1. Discovery Phase (discover_operations.py):

   • Scans all Python files in forge/forge/op/ directory
   • Identifies operation functions (functions starting with uppercase letters)
   • Uses Python AST to extract function signatures, docstrings, and parameters
   • Infers categories from file names (e.g., eltwise_unary.py → "Elementwise Operations")

2. Conversion Phase (generate_ops_docs.py):

   • Converts discovered operations to documentation format
   • Cleans up descriptions and removes redundant type prefixes
   • Handles edge cases (missing docstrings, incorrect descriptions)
   • Infers better descriptions from function names when docstrings are inadequate

3. Generation Phase (generate_ops_docs.py):

   • Creates individual markdown pages for each operation
   • Generates categorized index page
   • Formats in PyTorch-style documentation format

Usage

To regenerate documentation after adding new operations:

python scripts/generate_ops_docs.py

New operations are automatically discovered if they:

• Are defined in forge/forge/op/*.py files
• Have function names starting with uppercase letters
• Include proper docstrings (preferably NumPy-style with Parameters and Returns sections)

Documentation Structure

Index Page (operations.md)

• Operations organized by categories:
  • Elementwise Operations
  • Convolution Functions
  • Pooling Functions
  • Linear Functions
  • Reduction Operations
  • Tensor Manipulation
  • Normalization Functions
  • Resize Operations
  • Embedding Functions
  • Memory Operations
  • Creation Operations
  • Other Operations

Individual Operation Pages

Each page includes:

• Function Signature: Extracted from Python source code with type annotations
• Parameters: From docstrings with types, descriptions, and default values
• Returns: Return value descriptions with types
• Categories: Automatically inferred from file names

Current Coverage

Automatically discovered 78 operations from the codebase, including:

• Elementwise: Abs, Add, Subtract, Multiply, Divide, Relu, Sigmoid, Tanh, Gelu, etc.
• Convolution: Conv2d, Conv2dTranspose
• Pooling: MaxPool1d, MaxPool2d, AvgPool1d, AvgPool2d
• Linear: Matmul
• Reduction: ReduceSum, ReduceAvg, ReduceMax, Argmax
• Tensor Manipulation: Reshape, Transpose, Broadcast, Pad, etc.
• Normalization: Softmax, LogSoftmax, Layernorm, Batchnorm, Dropout
• And many more...

Testing

• ✅ Successfully discovers 78 operations from codebase
• ✅ Generates all operation pages correctly
• ✅ Creates categorized index page
• ✅ Handles edge cases (missing docstrings, incorrect descriptions)
• ✅ Function signatures match actual Python code
• ✅ All links working correctly
• ✅ Integrated with mdBook documentation system

This implementation provides:

• ✅ Automatic discovery of operations from forge/forge/op/*
• ✅ Structured operation list similar to PyTorch
• ✅ Individual pages for each operation with all details
• ✅ Root page with operations split by types with short descriptions
• ✅ Automatic discovery of new operations (no manual entry needed)

Related Issue

Closed: #59

[nvukobratTT]

@samiyadev786 thanks for your contribution! We'll review it asap! :))

@chandrasekaranpradeep, can you do an initial overview of the change? Have in mind that operation documentation should be generated based on existing ops + automatically discover when new are added and generate docs based on them as well. I.e., referncing this and similar files:

• tenstorrent/tt-forge-fe/forge/forge/op/*

[chandrasekaranpradeep]

@samiyadev786 thanks for your contribution! We'll review it asap! :))

@chandrasekaranpradeep, can you do an initial overview of the change? Have in mind that operation documentation should be generated based on existing ops + automatically discover when new are added and generate docs based on them as well. I.e., referncing this and similar files:

• tenstorrent/tt-forge-fe/forge/forge/op/*

Sure Nikola will review it

[samiyadev786]

@chandrasekaranpradeep Please review this PR.

[chandrasekaranpradeep]

@chandrasekaranpradeep Please review this PR.

I saw doc generation implementation changes from ttir to forge operation. So waiting for you to complete. will review this PR within today or tomorrow

[samiyadev786]

@chandrasekaranpradeep Please review this PR.

I saw doc generation implementation changes from ttir to forge operation. So waiting for you to complete. will review this PR within today or tomorrow

@chandrasekaranpradeep Thank you! The implementation is already complete on my end. Please feel free to review the PR when you’re available.