Ecosyste.ms: Awesome
An open API service indexing awesome lists of open source software.
https://github.com/radenmuaz/slope-ad
A small automatic differentiation engine, supporting higher-order derivatives
https://github.com/radenmuaz/slope-ad
array autograd automatic-differentiation cuda gradient iree jvp machine-learning metal mlir onnx onnxruntime tensor vjp
Last synced: 2 months ago
JSON representation
A small automatic differentiation engine, supporting higher-order derivatives
- Host: GitHub
- URL: https://github.com/radenmuaz/slope-ad
- Owner: radenmuaz
- License: mit
- Created: 2023-02-24T14:51:36.000Z (almost 2 years ago)
- Default Branch: main
- Last Pushed: 2024-06-09T11:03:06.000Z (8 months ago)
- Last Synced: 2024-11-15T21:44:15.390Z (3 months ago)
- Topics: array, autograd, automatic-differentiation, cuda, gradient, iree, jvp, machine-learning, metal, mlir, onnx, onnxruntime, tensor, vjp
- Language: Python
- Homepage: https://radenmuaz.github.io/slope-ad/
- Size: 4.54 MB
- Stars: 5
- Watchers: 2
- Forks: 1
- Open Issues: 0
-
Metadata Files:
- Readme: README.md
- License: LICENSE
Awesome Lists containing this project
README
# SlopeAD
Slope is a small automatic differentation (AD) engine, focused on machine learning (ML), supporting forward, reverse and higher-order AD.
This project is designed to be a **small**, **hackable** and **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/).
Example:
```python
import slope
def f(x):
y = x * 2.0
return y.sum()
x = slope.tensor([1.,2.,3.])
gf_x = slope.grad(f)(x)
print(f"{gf_x=}")
```
```
gf_x=
```
# Install
```
pip install slope-ad
```
or latest from main branch:
```
git clone https://github.com/radenmuaz/slope-ad
cd slope
pip install -e .
```
or you can just copy `src/slope` to your projects.
# Features
1. Small (?)
- <3000 lines of core code [slope/core.py](./src/slope/core.py), after formatted with `black src --line-length 140`
2. 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, where code is compiled to these supported backends running on either CPU, CUDA and Metal:
- [ONNX Runtime](https://onnxruntime.ai/) (ONNX graph)
- [OpenXLA IREE](https://iree.dev/) (StableHLO MLIR)
- NumPy (Python code)
4. 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)
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)
- 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.
- Procedures are exposed with `Tensor.procedure_name(*args)` syntax.
6. Extensible
- Add new backend by defining implementation translations [slope/backends](./src/slope/backends)
- Define new modules with NN module [slope/nn.py](./src/slope/nn.py)
# Docs
Docs are available online at [radenmuaz.github.io/slope-ad](https://radenmuaz.github.io/slope-ad)
API reference: [radenmuaz.github.io/slope-ad/api](https://radenmuaz.github.io/slope-ad/api)
## Tutorials
[Quickstart](./docs/tutorials/quickstart.md): How Tensors work, how to write and jit compile functions and train something.
[NN Training](./docs/tutorials/nn_training.md): Train MLP on MNIST with slope.nn module
[Internals Walkthrough](./docs/tutorials/internals_walkthrough.md): Understand the core of SlopeAD (hint: like JAX). Useful if you want to start contributing to SlopeAD
[Extending SlopeAD](./docs/tutorials/internals_walkthrough.md): Add new backend, operators, procedures. Modify the core functions.
## API reference
# Contributing
Open a PR, things on the roadmap below need to be done.
# Roadmap
- Docs
- Symbolic shape inference
- Dynamic shape jit
- Optimizer filter frozen params
- vmap vjp and jvp to compute jacobian and hessian
- iree backend currently has fixed seed random, implement threefry and JAX-like random
- make things fast
- llama (gpt) training
- whisper inference
- core tests, operators tests on all Trace types