diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml new file mode 100644 index 0000000000..b8074363c6 --- /dev/null +++ b/.github/workflows/docs.yml @@ -0,0 +1,29 @@ +name: Deploy Docs +on: + push: + branches: + - master + - mkdocs +permissions: + contents: write +jobs: + deploy: + runs-on: ubuntu-latest + steps: + - uses: actions/checkout@v4 + - name: Configure Git Credentials + run: | + git config user.name github-actions[bot] + git config user.email 41898282+github-actions[bot]@users.noreply.github.com + - uses: actions/setup-python@v5 + with: + python-version: 3.x + - run: echo "cache_id=$(date --utc '+%V')" >> $GITHUB_ENV + - uses: actions/cache@v4 + with: + key: mkdocs-material-${{ env.cache_id }} + path: .cache + restore-keys: | + mkdocs-material- + - run: pip install -e .[docs] + - run: mkdocs gh-deploy --force \ No newline at end of file diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml index 5412937a70..decbc85af1 100644 --- a/.github/workflows/test.yml +++ b/.github/workflows/test.yml @@ -91,7 +91,7 @@ jobs: run: python -m mypy --strict-equality - name: Test Docs run: | - python docs/abstractions2.py + python docs-legacy/abstractions2.py - name: Test Quickstart run: awk '/```python/{flag=1;next}/```/{flag=0}flag' docs/quickstart.md > quickstart.py && PYTHONPATH=. python quickstart.py - name: Fuzz Test symbolic diff --git a/.gitignore b/.gitignore index e428b9091c..6751355db9 100644 --- a/.gitignore +++ b/.gitignore @@ -51,3 +51,4 @@ quickstart.py .hypothesis weights *.lprof +site/ diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml index 13a0324240..f254c7c73b 100644 --- a/.pre-commit-config.yaml +++ b/.pre-commit-config.yaml @@ -21,7 +21,7 @@ repos: pass_filenames: false - id: docs2 name: docs2 - entry: python3 docs/abstractions2.py + entry: python3 docs-legacy/abstractions2.py language: system always_run: true pass_filenames: false diff --git a/docs/DESIGNv2.md b/docs-legacy/DESIGNv2.md similarity index 100% rename from docs/DESIGNv2.md rename to docs-legacy/DESIGNv2.md diff --git a/docs/OVERVIEW.md b/docs-legacy/OVERVIEW.md similarity index 100% rename from docs/OVERVIEW.md rename to docs-legacy/OVERVIEW.md diff --git a/docs/README.md b/docs-legacy/README.md similarity index 100% rename from docs/README.md rename to docs-legacy/README.md diff --git a/docs/abstractions2.py b/docs-legacy/abstractions2.py similarity index 100% rename from docs/abstractions2.py rename to docs-legacy/abstractions2.py diff --git a/docs/abstractions3.py b/docs-legacy/abstractions3.py similarity index 100% rename from docs/abstractions3.py rename to docs-legacy/abstractions3.py diff --git a/docs/adding_new_accelerators.md b/docs-legacy/adding_new_accelerators.md similarity index 100% rename from docs/adding_new_accelerators.md rename to docs-legacy/adding_new_accelerators.md diff --git a/docs/env_vars.md b/docs-legacy/env_vars.md similarity index 100% rename from docs/env_vars.md rename to docs-legacy/env_vars.md diff --git a/docs/linearizer_v2.md b/docs-legacy/linearizer_v2.md similarity index 100% rename from docs/linearizer_v2.md rename to docs-legacy/linearizer_v2.md diff --git a/docs/logo_tiny_dark.svg b/docs-legacy/logo_tiny_dark.svg similarity index 100% rename from docs/logo_tiny_dark.svg rename to docs-legacy/logo_tiny_dark.svg diff --git a/docs/logo_tiny_light.svg b/docs-legacy/logo_tiny_light.svg similarity index 100% rename from docs/logo_tiny_light.svg rename to docs-legacy/logo_tiny_light.svg diff --git a/docs/reshape_without_symbolic.md b/docs-legacy/reshape_without_symbolic.md similarity index 100% rename from docs/reshape_without_symbolic.md rename to docs-legacy/reshape_without_symbolic.md diff --git a/docs/developer.md b/docs/developer.md new file mode 100644 index 0000000000..762d251db5 --- /dev/null +++ b/docs/developer.md @@ -0,0 +1,7 @@ +## Frontend + +Everything in [Tensor](tensor.md) is syntactic sugar around [function.py](function.md), where the forwards and backwards passes are implemented for the different ops. That goes on to construct a graph of + +::: tinygrad.lazy.LazyBuffer + options: + show_source: false diff --git a/docs/dtypes.md b/docs/dtypes.md new file mode 100644 index 0000000000..62a631d7a4 --- /dev/null +++ b/docs/dtypes.md @@ -0,0 +1,3 @@ +::: tinygrad.dtypes + options: + members: true \ No newline at end of file diff --git a/docs/function.md b/docs/function.md new file mode 100644 index 0000000000..40ca5730d6 --- /dev/null +++ b/docs/function.md @@ -0,0 +1,7 @@ + + +::: tinygrad.function + options: + members: true + inherited_members: false + show_source: false diff --git a/docs/index.md b/docs/index.md new file mode 100644 index 0000000000..381bb060db --- /dev/null +++ b/docs/index.md @@ -0,0 +1,37 @@ +Welcome to the docs for tinygrad. This page is for users of the tinygrad library. We also have [developer docs](developer.md) + +tinygrad is not 1.0 yet, but it will be soon. The API has been pretty stable for a while. + +## tinygrad Usage + +The main class you will interact with is [Tensor](tensor.md). It functions very similarly to PyTorch, but has a bit more of a functional style. tinygrad supports [many datatypes](dtypes.md). All operations in tinygrad are lazy, meaning they won't do anything until you realize. + +* tinygrad has a built in [neural network library](nn.md) with some classes, optimizers, and load/save state management. +* tinygrad has a JIT to make things fast. Decorate your pure function with `TinyJit` +* tinygrad has amazing support for multiple GPUs, allowing you to shard your Tensors with `Tensor.shard` + +To understand what training looks like in tinygrad, you should read `beautiful_mnist.py` + +We have a [quickstart guide](quickstart.md) and a [showcase](showcase.md) + +## Differences from PyTorch + +If you are migrating from PyTorch, welcome. We hope you will find tinygrad both familiar and somehow more "correct feeling" + +### tinygrad doesn't have nn.Module + +There's nothing special about a "Module" class in tinygrad, it's just a normal class. `get_parameter` + +### tinygrad is functional + + + +In tinygrad, you can do `x.conv2d(w, b)` or `x.sparse_categorical_cross_entropy(y)` + +### tinygrad is lazy + +When you do `a+b` in tinygrad, nothing happens. + +### tinygrad requires @TinyJIT to be fast + +PyTorch spends a lot of development effort to make dispatch very fast. tinygrad doesn't. We have a simple decorator that will replay the kernels used in the decorated function. \ No newline at end of file diff --git a/docs/nn.md b/docs/nn.md new file mode 100644 index 0000000000..25aa46ee51 --- /dev/null +++ b/docs/nn.md @@ -0,0 +1,17 @@ +## Neural Network classes + +::: tinygrad.nn + options: + members: true + +## Optimizers + +::: tinygrad.nn.optim + options: + members: true + +## Load/Save + +::: tinygrad.nn.state + options: + members: true \ No newline at end of file diff --git a/docs/showcase.md b/docs/showcase.md index d0b2d4630e..04b18b7be4 100644 --- a/docs/showcase.md +++ b/docs/showcase.md @@ -19,7 +19,7 @@ python3 examples/efficientnet.py webcam Take a look at [yolov8.py](/examples/yolov8.py). -![yolov8 by tinygrad](/docs/showcase/yolov8_showcase_image.png) +![yolov8 by tinygrad](showcase/yolov8_showcase_image.png) ## Audio @@ -37,7 +37,7 @@ SMALL=1 python3 examples/whisper.py Take a look at [mnist_gan.py](/examples/mnist_gan.py). -![mnist gan by tinygrad](/docs/showcase/mnist_by_tinygrad.jpg) +![mnist gan by tinygrad](showcase/mnist_by_tinygrad.jpg) ### Stable Diffusion @@ -45,7 +45,7 @@ Take a look at [mnist_gan.py](/examples/mnist_gan.py). python3 examples/stable_diffusion.py ``` -![a horse sized cat eating a bagel](/docs/showcase/stable_diffusion_by_tinygrad.jpg) +![a horse sized cat eating a bagel](showcase/stable_diffusion_by_tinygrad.jpg) *"a horse sized cat eating a bagel"* diff --git a/docs/tensor.md b/docs/tensor.md new file mode 100644 index 0000000000..e2378336c4 --- /dev/null +++ b/docs/tensor.md @@ -0,0 +1,167 @@ +::: tinygrad.Tensor + options: + heading_level: 2 + members: false + show_source: false + +## Properties + +::: tinygrad.Tensor.shape +::: tinygrad.Tensor.dtype +::: tinygrad.Tensor.device + +## tinygrad ops + +::: tinygrad.Tensor.corealize +::: tinygrad.Tensor.realize +::: tinygrad.Tensor.replace +::: tinygrad.Tensor.assign +::: tinygrad.Tensor.contiguous +::: tinygrad.Tensor.contiguous_backward + +## Creation (basic) + +::: tinygrad.Tensor.empty +::: tinygrad.Tensor.zeros +::: tinygrad.Tensor.ones +::: tinygrad.Tensor.full +::: tinygrad.Tensor.arange +::: tinygrad.Tensor.eye + +## Creation (random) + +::: tinygrad.Tensor.rand +::: tinygrad.Tensor.randn +::: tinygrad.Tensor.normal +::: tinygrad.Tensor.uniform +::: tinygrad.Tensor.scaled_uniform +::: tinygrad.Tensor.glorot_uniform +::: tinygrad.Tensor.kaiming_uniform +::: tinygrad.Tensor.kaiming_normal + +## Movement (low level) + +::: tinygrad.Tensor.reshape +::: tinygrad.Tensor.expand +::: tinygrad.Tensor.permute +::: tinygrad.Tensor.flip +::: tinygrad.Tensor.shrink +::: tinygrad.Tensor.pad + +## Movement (high level) + +::: tinygrad.Tensor.__getitem__ +::: tinygrad.Tensor.slice +::: tinygrad.Tensor.gather +::: tinygrad.Tensor.cat +::: tinygrad.Tensor.stack +::: tinygrad.Tensor.repeat +::: tinygrad.Tensor.split +::: tinygrad.Tensor.chunk +::: tinygrad.Tensor.squeeze +::: tinygrad.Tensor.unsqueeze +::: tinygrad.Tensor.pad2d +::: tinygrad.Tensor.transpose +::: tinygrad.Tensor.flatten +::: tinygrad.Tensor.unflatten + +## Reduce + +::: tinygrad.Tensor.sum +::: tinygrad.Tensor.max +::: tinygrad.Tensor.min +::: tinygrad.Tensor.mean +::: tinygrad.Tensor.var +::: tinygrad.Tensor.std +::: tinygrad.Tensor.softmax +::: tinygrad.Tensor.log_softmax +::: tinygrad.Tensor.argmax + +## Processing + +::: tinygrad.Tensor.conv2d +::: tinygrad.Tensor.dot +::: tinygrad.Tensor.matmul +::: tinygrad.Tensor.einsum +::: tinygrad.Tensor.cumsum +::: tinygrad.Tensor.triu +::: tinygrad.Tensor.tril +::: tinygrad.Tensor.avg_pool2d +::: tinygrad.Tensor.max_pool2d +::: tinygrad.Tensor.conv_transpose2d + +## Unary Ops (math) + +::: tinygrad.Tensor.logical_not +::: tinygrad.Tensor.neg +::: tinygrad.Tensor.log +::: tinygrad.Tensor.log2 +::: tinygrad.Tensor.exp +::: tinygrad.Tensor.exp2 +::: tinygrad.Tensor.trunc +::: tinygrad.Tensor.ceil +::: tinygrad.Tensor.floor +::: tinygrad.Tensor.round +::: tinygrad.Tensor.lerp +::: tinygrad.Tensor.square +::: tinygrad.Tensor.clip +::: tinygrad.Tensor.abs +::: tinygrad.Tensor.sign +::: tinygrad.Tensor.reciprocal + +## Unary Ops (activation) + +::: tinygrad.Tensor.relu +::: tinygrad.Tensor.sigmoid +::: tinygrad.Tensor.elu +::: tinygrad.Tensor.celu +::: tinygrad.Tensor.swish +::: tinygrad.Tensor.silu +::: tinygrad.Tensor.relu6 +::: tinygrad.Tensor.hardswish +::: tinygrad.Tensor.tanh +::: tinygrad.Tensor.sinh +::: tinygrad.Tensor.cosh +::: tinygrad.Tensor.atanh +::: tinygrad.Tensor.asinh +::: tinygrad.Tensor.acosh +::: tinygrad.Tensor.hardtanh +::: tinygrad.Tensor.gelu +::: tinygrad.Tensor.quick_gelu +::: tinygrad.Tensor.leakyrelu +::: tinygrad.Tensor.mish +::: tinygrad.Tensor.softplus +::: tinygrad.Tensor.softsign + +## Elementwise Ops (broadcasted) + +::: tinygrad.Tensor.add +::: tinygrad.Tensor.sub +::: tinygrad.Tensor.mul +::: tinygrad.Tensor.div +::: tinygrad.Tensor.xor +::: tinygrad.Tensor.pow +::: tinygrad.Tensor.maximum +::: tinygrad.Tensor.minimum +::: tinygrad.Tensor.where + +## Neural Network Ops (functional) + +::: tinygrad.Tensor.linear +::: tinygrad.Tensor.sequential +::: tinygrad.Tensor.layernorm +::: tinygrad.Tensor.batchnorm +::: tinygrad.Tensor.dropout +::: tinygrad.Tensor.one_hot +::: tinygrad.Tensor.scaled_dot_product_attention +::: tinygrad.Tensor.binary_crossentropy +::: tinygrad.Tensor.binary_crossentropy_logits +::: tinygrad.Tensor.sparse_categorical_crossentropy + +## Casting Ops + +::: tinygrad.Tensor.cast +::: tinygrad.Tensor.bitcast +::: tinygrad.Tensor.float +::: tinygrad.Tensor.half + diff --git a/mkdocs.yml b/mkdocs.yml new file mode 100644 index 0000000000..c553af902d --- /dev/null +++ b/mkdocs.yml @@ -0,0 +1,98 @@ +# pip install mkdocs mkdocs-material mkdocstrings[python] +site_name: tinygrad docs +site_url: https://docs.tinygrad.org/ +nav: +- Home: index.md +- Tensor: tensor.md +- dtypes: dtypes.md +- Neural Networks: nn.md +- Quickstart: quickstart.md +- Showcase: showcase.md +- Developer: developer.md +- Function: function.md +#- tinygrad: reference/ + +#extra_css: +#- css/tinygrad.css + +markdown_extensions: +- attr_list +- admonition +- callouts +- footnotes +- pymdownx.details +- pymdownx.emoji: + emoji_index: !!python/name:material.extensions.emoji.twemoji + emoji_generator: !!python/name:material.extensions.emoji.to_svg +- pymdownx.highlight: + pygments_lang_class: true +- pymdownx.inlinehilite: + style_plain_text: python +- pymdownx.magiclink +- pymdownx.snippets: + base_path: [!relative $config_dir] + check_paths: true +- pymdownx.superfences +- pymdownx.tabbed: + alternate_style: true + slugify: !!python/object/apply:pymdownx.slugs.slugify + kwds: + case: lower +- pymdownx.tasklist: + custom_checkbox: true +- pymdownx.tilde +- toc: + permalink: "ยค" + +theme: + name: material + features: + - announce.dismiss + - content.action.edit + - content.action.view + - content.code.annotate + - content.code.copy + - content.tooltips + - navigation.footer + - navigation.indexes + - navigation.sections + - navigation.tabs + - navigation.tabs.sticky + - navigation.top + - search.highlight + - search.suggest + - toc.follow + palette: + scheme: slate + primary: black + accent: lime +plugins: +- search +- mkdocstrings: + handlers: + python: + import: + - https://docs.python.org/3/objects.inv + paths: [tinygrad] + options: + docstring_options: + ignore_init_summary: true + docstring_section_style: list + filters: ["!^_"] + heading_level: 3 + inherited_members: false + merge_init_into_class: true + separate_signature: true + show_root_heading: true + show_root_full_path: false + show_signature_annotations: true + show_symbol_type_heading: true + show_symbol_type_toc: true + show_source: true + signature_crossrefs: true + summary: true +#- gen-files: +# scripts: +# - docs/gen_ref_pages.py +#- literate-nav: +# nav_file: SUMMARY.md \ No newline at end of file diff --git a/ruff.toml b/ruff.toml index 8bfcd3b8fc..bccdecbb93 100644 --- a/ruff.toml +++ b/ruff.toml @@ -28,6 +28,7 @@ line-length = 150 exclude = [ "disassemblers/", "docs/", + "docs-legacy/", "examples/", "extra/", "openpilot/", diff --git a/setup.py b/setup.py index 1180744f49..5a41e55f61 100644 --- a/setup.py +++ b/setup.py @@ -53,6 +53,11 @@ setup(name='tinygrad', "networkx", "hypothesis", ], + 'docs': [ + "mkdocs-material", + "mkdocstrings[python]", + "markdown-callouts", + ], 'testing_tf': [ "tensorflow==2.15.1", "tensorflow_addons", diff --git a/tinygrad/function.py b/tinygrad/function.py index 9c854fc48a..15bbc2800b 100644 --- a/tinygrad/function.py +++ b/tinygrad/function.py @@ -1,3 +1,4 @@ +"""This is where the forwards and backwards passes live.""" import math from typing import Tuple, Optional from tinygrad.helpers import argsort diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 35be3fe728..1352254da5 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -70,6 +70,7 @@ def _pad_left(*shps:Tuple[sint, ...], v=1): return tuple((v,) * (max(len(i_) for def broadcast_shape(*shps:Tuple[sint, ...]): return tuple(0 if any(sh_ == 0 for sh_ in sh) else max(sh) for sh in zip(*_pad_left(*shps))) class Tensor: + """A `Tensor` is a multi-dimensional matrix containing elements of a single data type.""" __slots__ = "lazydata", "requires_grad", "grad", "_ctx" __deletable__ = ('_ctx',) training: ClassVar[bool] = False @@ -836,7 +837,6 @@ class Tensor: def round(self: Tensor) -> Tensor: return ((self > 0) == ((b := self.cast(dtypes.int32) / 2.0).cast(dtypes.int32) == b)).where((self - 0.5).ceil(), (self + 0.5).floor()) def lerp(self, end: Tensor, weight: Union[Tensor, float]) -> Tensor: return self + (end - self) * weight - def square(self): return self*self def clip(self, min_, max_): return self.maximum(min_).minimum(max_) def abs(self): return self.relu() + (-self).relu()