diff --git a/README.md b/README.md index d6f03c9..00b8bfd 100644 --- a/README.md +++ b/README.md @@ -44,11 +44,21 @@ pip install -e . Or you can just copy `src/slope` to your projects. -# Docs +# Usage -Docs are in [radenmuaz.github.io/slope-ad](https://radenmuaz.github.io/slope-ad) +## Tutorials -For a quick overview on how to train and backprop something, go to the [Quickstart Tutorial](./docs/tutorials/quickstart.md) +[Quickstart Tutorial](./docs/tutorials/quickstart.md): How Tensors work, how to write and jit compile functions and train something + +[NN Training](./docs/tutorials/quickstart.md): Train a MLP on MNIST and ResNet on CIFAR-10 + +[Internals Walkthrough](./docs/tutorials/internals_walkthrough.md): Understand the core of SlopeAD (hint: like JAX) + +[Extending SlopeAD](./docs/tutorials/internals_walkthrough.md): Add new backend, operators, procedures. Modify the core functions. + +## Docs and API reference + +Docs are available online at [radenmuaz.github.io/slope-ad](https://radenmuaz.github.io/slope-ad) # Features @@ -76,11 +86,11 @@ For a quick overview on how to train and backprop something, go to the [Quicksta - Init: `full arange random_normal random_uniform` - GeneralReduce: `matmul conv gather_nd scatter_nd` - Composite operators system with "procedures" [slope/procedures.py](./src/slope/procedures.py) - - Procedures are functions containing calls to operators, exposed with `Tensor.procedure_name(*args)` syntax. - - Useful for definitions like: + - For defining Tensor functions composed with core operators, e.g. - `x.cos()`, where `def cos(x): return (math.pi/2 - x).sin()` - `x.conv_transpose(w)`: where `def conv_transpose(x, w): ...` is a very long function. - - An operator can be directly implemented as code translation to backend, or _fallback_ to a procedure, e.g. there is `conv` procedure in case the backend has no implementation for it. + - Procedures are functions containing calls to operators, exposed with `Tensor.procedure_name(*args)` syntax. + 6. Extensible - Add new backend by defining implementation translations [slope/backends](./src/slope/backends) diff --git a/docs/index.md b/docs/index.md index 05be8eb..0da06ab 100644 --- a/docs/index.md +++ b/docs/index.md @@ -2,7 +2,7 @@ Slope is a small automatic differentation (AD) engine, focused on machine learning (ML) -This project is designed to be a hackable, educational AD engine focused on ML, yet able to do things end-to-end training to deployment, instead of just some simple toy examples. +This project is designed to be a **hackable**, **educational** AD engine focused on ML, yet able to do things **end-to-end from training to deployment**, instead of just some simple toy examples. Tensor semantics are similar to Pytorch, functional API is similar to [JAX](https://github.com/google/jax), tensor operators code is heavily derived from [tinygrad](https://tinygrad.org/). @@ -25,6 +25,26 @@ shape=(3,), dtype=float32, device='cpu:0'> ``` +# Features + +1. Functional API for forward-mode, reverse-mode, and higher-order AD, like in JAX: + - `grad vjp jvp jit vmap` + - `register_node tree_flatten tree_unflatten` + +3. Just-in-time compilation, code is compiled to one of these backends: [ONNX Runtime](https://onnxruntime.ai/), [OpenXLA IREE](https://iree.dev/) and NumPy. + +4. Training and inference, like [MLP on MNIST](examples/nn/mnist_mlp.py), [ResNet on CIFAR-10](examples/nn/cifar_resnet.py) and [export jitted function](examples/simple/export.py). + +5. Small (?), less than 3000 lines of core code [slope/core.py](./src/slope/core.py), after `black src --line-length 140` + +6. Operators and procedures system + - 33 core operators defined in [slope/operators.py](./src/slope/operators.py) + - `exp log sin sqrt invert cast stop_gradient add mul sub div pow equal less greater maximum sum max reshape expand permute slice pad flip cat full arange random_normal random_uniform matmul conv gather_nd scatter_nd` + - Composite operators system with "procedures" [slope/procedures.py](./src/slope/procedures.py), for defining functions like `cos`, `conv` and `avgpool2d` as functions calling core operators. + +7. Extensible, by writing new backend by defining implementation translations [slope/backends](./src/slope/backends), and adding more modules using [slope/nn.py](./src/slope/nn.py) + + # Install ## Stable release @@ -43,45 +63,23 @@ pip install -e . Or you can just copy `src/slope` to your projects. -# Docs +# Usage -Head over to [quickstart](./tutorials/quickstart.md) +## Tutorials -# Features +[Quickstart Tutorial](./docs/tutorials/quickstart.md): How Tensors work, how to write and jit compile functions and train something -1. Forward-mode, reverse-mode, and higher-order AD. +[NN Training](./docs/tutorials/quickstart.md): Train a MLP on MNIST and ResNet on CIFAR-10 -2. Just-in-time compilation, with interchangeable backends, supporting CPU, CUDA and Metal: - - [IREE](https://iree.dev/) (StableHLO MLIR) - - [ONNX Runtime](https://onnxruntime.ai/) (ONNX) - - NumPy (CPU-only) +[Internals Walkthrough](./docs/tutorials/internals_walkthrough.md): Understand the core of SlopeAD (hint: like JAX) -3. Training and inference, examples: - - [MLP on MNIST](examples/nn/mnist_mlp.py) - - [ResNet on CIFAR-10](examples/nn/cifar_resnet.py) - - [Export jitted function](examples/simple/export.py) +[Extending SlopeAD](./docs/tutorials/internals_walkthrough.md): Add new backend, operators, procedures. Modify the core functions. + +## Docs and API reference + +Docs are available online at [radenmuaz.github.io/slope-ad](https://radenmuaz.github.io/slope-ad) -4. Small (?) - - <3000 lines of core code [slope/core.py](./src/slope/core.py), after run with `black src --line-length 140` -5. Operators and procedures system - - 33 core operators defined in [slope/operators.py](./src/slope/operators.py) - - Unary: `exp log sin sqrt invert cast stop_gradient` - - Binary: `add mul sub div pow equal less greater maximum` - - Reduce: `sum max` - - Shape: `reshape expand permute slice pad flip cat` - - Init: `full arange random_normal random_uniform` - - GeneralReduce: `matmul conv gather_nd scatter_nd` - - Composite operators system with "procedures" [slope/procedures.py](./src/slope/procedures.py) - - Procedures are functions containing calls to operators, exposed with `Tensor.procedure_name(*args)` syntax. - - Useful for definitions like: - - `x.cos()`, where `def cos(x): return (math.pi/2 - x).sin()` - - `x.conv_transpose(w)`: where `def conv_transpose(x, w): ...` is a very long function. - - An operator can be directly implemented as code translation to backend, or _fallback_ to a procedure, e.g. there is `conv` procedure in case the backend has no implementation for it. - -6. Extensible - - Add new backend by defining implementation translations [slope/backends](./src/slope/backends) - - NN module [slope/nn.py](./src/slope/nn.py) diff --git a/docs/tutorials/ad_vmap_tree.md b/docs/tutorials/ad_vmap_tree.md deleted file mode 100644 index e69de29..0000000