TinyGrad Codebase Explained-ish

👋 Hello lovely readers! Let’s dive into TinyGrad’s codebase structure and key components as of commit #1efb1188d8.

This is a high-level overview sharing my learnings as I explore the codebase 🤓

Note: Code outside the core tinygrad/ directory is not extensively tested. Unless the code is broken, it’s advisable not to modify it.

Core TinyGrad Directory

Let’s explore the heart of TinyGrad - where the core implementation lives:

Core Files

# Core Python files
tensor.py    # Tensor implementation with autograd
ops.py       # Base operations and math traits
device.py    # Hardware abstraction
dtype.py     # Data type system
function.py  # Function transformations
gradient.py  # Automatic differentiation
multi.py     # Multi-device operations
helpers.py   # Utility functions
__init__.py  # Package initialization
py.typed     # Type checking marker

Key Directories

• engine/: Memory management, operation scheduling, JIT compilation, Graph realization

• runtime/: Hardware-specific operations (CUDA, Metal, WebGPU), platform-specific implementations, support for hardware compilation, auto-generated hardware bindings, graph implementations for different backends

• codegen/: Kernel generation, operation linearization, graph optimization, lowering operations

• shape/: Shape tracking, view operations

• nn/: Optimizers, neural network state, datasets

• renderer/: LLVM IR generation, PTX code generation, WGSL (WebGPU Shading Language) generation, C-style code generation

• viz/: Debugging, visualization, profiling, and graph tracking. Components: serve.py, index.html, perfetto.html, README, fetch_assets.sh, assets/

Documentation (docs.tinygrad.org)

The docs folder contains all TinyGrad documentation. For a quick overview, refer to tinygrad_intro.pdf.

Documentation Tools

• serve_docs.sh: This is a bash script that runs the mkdocs serve command with the -w (watch) flag, pointing to the tinygrad/ directory. Its purpose is to serve documentation locally while watching the tinygrad directory for changes.

• mkdocs.yml: This is a MkDocs configuration file (mkdocs.yml) that sets up the documentation website for tinygrad, defining the site structure, navigation menu, theme settings, and various Markdown extensions. The file configures a Material theme with dark/light mode support and includes plugins for Python documentation generation (mkdocstrings), search functionality, and markdown execution.

Development and Testing Tools

TinyGrad takes code quality seriously. Here are the key tools that maintain the project’s high standards:

Code Quality Configuration

• ruff.toml: This is a ruff.toml configuration file that sets up Python linting rules with specific settings like indent width of 2 spaces, Python 3.10 as the target version, and a line length limit of 150 characters. The file selects various linting rules across categories like Pyflakes, whitespace, indentation, and code style while excluding certain directories from linting checks.

• mypy.ini: This is a mypy.ini configuration file that sets up type checking rules for the tinygrad project. It specifies various type checking behaviors like ignoring missing imports, checking untyped definitions, and enabling warnings for unreachable code and redundant casts, with a note that the warn_unused_ignores option had to be commented out for compatibility between CI and OSX.

• eslint.config.mjs: This is an ESLint configuration file that exports default settings for code linting. It configures HTML file rules with a maximum line length of 150 characters, includes browser globals, and applies JavaScript recommended configurations using plugins for HTML and JavaScript.

• .pylintrc: This .pylintrc file is a configuration file for Pylint (a Python code analysis tool) that defines various rules and settings for code analysis, including disabled checks, naming conventions, and formatting requirements. It specifies settings across multiple categories like master controls, message handling, code design limits, and import rules, with notable configurations like a max line length of 150 characters and disabled checks for certain warning categories (C, R, and specific error codes).

• .pre-commit-config.yaml: This is a .pre-commit-config.yaml configuration file that defines several local hooks for code quality and testing in a Python project. The hooks run various checks including Ruff linting, MyPy type checking, PyLint, and multiple test suites (including tiny tests, GPU tests, and multi-device tests) that execute automatically before commits.

Test Folder

Quality assurance is built into TinyGrad’s DNA. Here’s how the project ensures reliability:

🧪 The test folder ensures code quality and correctness.

The test folder contains:

• Unit tests for core functionality
• Validation of primary operations
• Focus on tinygrad/ directory components

test_driven_development.sh

This Bash script:

• Runs a series of Python tests: test_tiny.py, test_test_uop_graph.py, test_ops.py, and test_linearizer.py
• Includes process replay functionality
• Enters an infinite loop to continuously run test_tiny.py
• Executes a process replay script if the test passes

setup.py

The project’s installation configuration:

• Version: 0.10.0
• License: MIT
• Description: “You like pytorch? You like micrograd? You love tinygrad! <3”
• Python 3.10+ requirement
• Optional dependency groups for:
  • LLVM integration
  • Testing tools
  • Documentation
  • TensorFlow compatibility

sz.py

A sophisticated analysis tool for Python codebases that:

• Generates statistics about line counts and token density
• Excludes docstrings and auto-generated files
• Enables comparison between codebases
• Shows changes in line counts and token density
• Formats results in tabular form by directory

Automation Scripts

• autogen_stubs.sh: This bash script generates Python bindings (stubs) for various hardware APIs and drivers (including OpenCL, CUDA, HIP, AMD, NVIDIA, etc.) by using clang2py to convert C/C++ headers to Python code. The script contains multiple functions to generate different bindings, applies necessary fixes and patches to the generated code, and places the output files in the tinygrad/runtime/autogen/ directory.

GitHub Workflows (.github/workflows)

Automation is crucial for maintaining quality. These GitHub workflows handle everything from testing to deployment:

• docs.yml: Handles documentation deployment to GitHub Pages when changes are pushed to the master branch.

• python-publish.yml: Manages Python package publishing to PyPI when a new release is created.

• szdiff.yml: Tracks line count differences in pull requests and checks if branches are up-to-date.

• test.yml: Runs comprehensive unit tests, linting checks, and various environment tests (OpenCL, CUDA, AMD) for code quality.

• benchmark.yml: Executes performance benchmarks across different platforms (Mac, NVIDIA) for various ML models and operations.

Key Points:

• All workflows are automated using GitHub Actions.
• Focus on testing, documentation, deployment, and performance benchmarking.
• Includes cross-platform testing (Mac, Linux, NVIDIA).
• Comprehensive coverage of ML operations and model testing.

Examples Folder

The best way to understand TinyGrad is to see it in action. These examples showcase its capabilities across different domains:

The examples folder contains working implementations of popular models!

The examples folder showcases TinyGrad’s capabilities:

• ASR: Whisper implementation
• Computer Vision: YOLOv8 and EfficientNet
• Language Models: LLaMA
• Machine Learning: Model training and benchmarking
• Neural network and machine learning examples
• Various other demonstrations and random stuff not covered in the blog post

Extra Folder

The extra folder contains supplementary code and resources that extend TinyGrad’s core functionality. These components are not essential for the basic operation of TinyGrad but provide useful tools and examples for users:

• Dataset loaders (e.g., MNIST)
• Additional utilities and tools
• Supplementary implementations like HIP (Heterogeneous-compute Interface for Portability) matrix multiplication in hip_matmul.py

Adel Aloui

Software Engineer with expertise in building innovative AI-driven solutions.

Tweet Share