diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
index 680848810b..41054519cc 100644
--- a/.github/workflows/test.yml
+++ b/.github/workflows/test.yml
@@ -55,6 +55,8 @@ jobs:
python-version: 3.8
- name: Install Dependencies
run: pip install -e '.[testing]' --extra-index-url https://download.pytorch.org/whl/cpu
+ - name: Test Docs
+ run: python docs/abstractions.py
- name: Run Pytest
run: python -m pytest -s -v -n=auto test/
diff --git a/.gitignore b/.gitignore
index 31745b13ec..ce863eb050 100644
--- a/.gitignore
+++ b/.gitignore
@@ -28,4 +28,5 @@ datasets/squad/
datasets/img_align_celeba*
datasets/open-images-v6-mlperf
datasets/kits/
+datasets/audio*
venv
diff --git a/.tokeignore b/.tokeignore
new file mode 100644
index 0000000000..cb7645b24d
--- /dev/null
+++ b/.tokeignore
@@ -0,0 +1,4 @@
+*
+!*/
+
+!tinygrad/**
diff --git a/README.md b/README.md
index 1620b777b1..3f56ee9b57 100644
--- a/README.md
+++ b/README.md
@@ -1,91 +1,57 @@
-
- 
-
+
---------------------------------------------------------------------
+[](https://tinygrad.org)
-
+tinygrad: For something between [PyTorch](https://github.com/pytorch/pytorch) and [karpathy/micrograd](https://github.com/karpathy/micrograd). Maintained by [tiny corp](https://tinygrad.org).
-[](https://discord.gg/ZjZadyC7PK)
+
-For something in between a [pytorch](https://github.com/pytorch/pytorch) and a [karpathy/micrograd](https://github.com/karpathy/micrograd)
+[Homepage](https://github.com/geohot/tinygrad) | [Documentation](/docs) | [Examples](/examples) | [Showcase](/docs/showcase.md) | [Discord](https://discord.gg/ZjZadyC7PK)
+
+
+
+[](https://github.com/geohot/tinygrad/stargazers)
+[](https://github.com/geohot/tinygrad/actions/workflows/test.yml)
+[](https://discord.gg/ZjZadyC7PK)
+[](https://github.com/geohot/tinygrad)
+
+
+
+---
This may not be the best deep learning framework, but it is a deep learning framework.
-The sub 1000 line core of it is in `tinygrad/`
+Due to its extreme simplicity, it aims to be the easiest framework to add new accelerators to, with support for both inference and training.
-Due to its extreme simplicity, it aims to be the easiest framework to add new accelerators to, with support for both inference and training. Support the simple basic ops, and you get SOTA [vision](https://arxiv.org/abs/1905.11946) `models/efficientnet.py` and [language](https://arxiv.org/abs/1706.03762) `models/transformer.py` models.
+Eventually, we will have a [tinygrad accelerator](https://geohot.github.io/blog/jekyll/update/2021/06/13/a-breakdown-of-ai-chip-companies.html), then tinygrad will be ***fast***. But, for now, it is slow.
-We are working on support for the Apple Neural Engine and the Google TPU in the `accel/` folder. Eventually, [we will build custom hardware](https://geohot.github.io/blog/jekyll/update/2021/06/13/a-breakdown-of-ai-chip-companies.html) for tinygrad, and it will be blindingly fast. Now, it is slow.
+## Features
-This project is maintained by [tiny corp](https://tinygrad.org/).
+### LLaMA and Stable Diffusion
-### Installation
+tinygrad can run [LLaMA](/docs/showcase.md#llama) and [Stable Diffusion](/docs/showcase.md#stable-diffusion)!
-```bash
-git clone https://github.com/geohot/tinygrad.git
-cd tinygrad
-python3 -m pip install -e .
-```
-
-### Contributing
-
-There's a lot of interest in tinygrad lately. Here's some guidelines for contributing:
-
-* Bugfixes are the best and always welcome! Like [this one](https://github.com/geohot/tinygrad/pull/421/files).
-* If you don't understand the code you are changing, don't change it!
-* All code golf PRs will be closed, but [conceptual cleanups](https://github.com/geohot/tinygrad/pull/372/files) are great.
-* Features are welcome. Though if you are adding a feature, you need to include tests.
-* Improving test coverage is great, with reliable non brittle tests.
-
-### Example
-
-```python
-from tinygrad.tensor import Tensor
-
-x = Tensor.eye(3, requires_grad=True)
-y = Tensor([[2.0,0,-2.0]], requires_grad=True)
-z = y.matmul(x).sum()
-z.backward()
-
-print(x.grad.numpy()) # dz/dx
-print(y.grad.numpy()) # dz/dy
-```
-
-### Same example in torch
-
-```python
-import torch
-
-x = torch.eye(3, requires_grad=True)
-y = torch.tensor([[2.0,0,-2.0]], requires_grad=True)
-z = y.matmul(x).sum()
-z.backward()
-
-print(x.grad) # dz/dx
-print(y.grad) # dz/dy
-```
-
-## Is tinygrad fast?
+### Laziness
Try a matmul. See how, despite the style, it is fused into one kernel with the power of laziness.
-```python
+```sh
DEBUG=3 OPTLOCAL=1 python3 -c "from tinygrad.tensor import Tensor;
-N = 1024; a, b = Tensor.randn(N, N), Tensor.randn(N, N);
+N = 1024; a, b = Tensor.rand(N, N), Tensor.rand(N, N);
c = (a.reshape(N, 1, N) * b.permute(1,0).reshape(1, N, N)).sum(axis=2);
print((c.numpy() - (a.numpy() @ b.numpy())).mean())"
```
-Change to `DEBUG=4` to see the generated code.
+And we can change `DEBUG` to `4` to see the generated code.
-## Neural networks?
+### Neural networks
-It turns out, a decent autograd tensor library is 90% of what you need for neural networks. Add an optimizer (SGD, Adam, AdamW implemented) from tinygrad.nn.optim, write some boilerplate minibatching code, and you have all you need.
+As it turns out, 90% of what you need for neural networks are a decent autograd/tensor library.
+Throw in an optimizer, a data loader, and some compute, and you have all you need.
-### Neural network example (from test/models/test_mnist.py)
+#### Neural network example (from test/models/test_mnist.py)
-```python
+```py
from tinygrad.tensor import Tensor
import tinygrad.nn.optim as optim
@@ -100,7 +66,7 @@ class TinyBobNet:
model = TinyBobNet()
optim = optim.SGD([model.l1, model.l2], lr=0.001)
-# ... and complete like pytorch, with (x,y) data
+# ... complete data loader here
out = model.forward(x)
loss = out.mul(y).mean()
@@ -109,114 +75,86 @@ loss.backward()
optim.step()
```
-## GPU and Accelerator Support
+## Accelerators
-tinygrad supports GPUs through PyOpenCL.
+tinygrad already supports numerous accelerators, including:
-```python
+- [x] CPU
+- [x] GPU (OpenCL)
+- [x] C Code (Clang)
+- [x] LLVM
+- [x] METAL
+- [x] CUDA
+- [x] Triton
+- [x] PyTorch
+
+And it is easy to add more! Your accelerator of choice only needs to support a total of 26 (optionally 27) low level ops.
+More information can be found in the [documentation for adding new accelerators](/docs/adding_new_accelerators.md).
+
+## Installation
+
+The current recommended way to install tinygrad is from source.
+
+### From source
+
+```sh
+git clone https://github.com/geohot/tinygrad.git
+cd tinygrad
+python3 -m pip install -e . # or `py3 -m pip install -e .` if you are on windows
+```
+Don't forget the `.` at the end!
+
+## Documentation
+
+Documentation along with a quick start guide can be found in the [docs/](/docs) directory.
+
+### Quick example comparing to PyTorch
+
+```py
from tinygrad.tensor import Tensor
-(Tensor.ones(4,4).gpu() + Tensor.ones(4,4).gpu()).cpu()
+
+x = Tensor.eye(3, requires_grad=True)
+y = Tensor([[2.0,0,-2.0]], requires_grad=True)
+z = y.matmul(x).sum()
+z.backward()
+
+print(x.grad.numpy()) # dz/dx
+print(y.grad.numpy()) # dz/dy
```
-### hlops (in tensor.py)
+The same thing but in PyTorch:
+```py
+import torch
-hlops are syntactic sugar around mlops. They support most things torch does.
+x = torch.eye(3, requires_grad=True)
+y = torch.tensor([[2.0,0,-2.0]], requires_grad=True)
+z = y.matmul(x).sum()
+z.backward()
-### mlops
-
-mlops are mid level ops. They understand derivatives. They are very simple.
-
-```
-Relu, Log, Exp, Sin # unary ops
-Sum, Max # reduce ops (with axis argument)
-Maximum, Add, Sub, Mul, Pow, Div, Equal # binary ops (no broadcasting, use expand)
-Expand, Reshape, Permute, Pad, Shrink, Flip # movement ops
+print(x.grad.numpy()) # dz/dx
+print(y.grad.numpy()) # dz/dy
```
-You no longer need to write mlops for a new accelerator
+## Contributing
-### Adding an accelerator (llops)
+There has been a lot of interest in tinygrad lately. Here are some basic guidelines for contributing:
-The autodiff stuff is all in mlops now so you can focus on the raw operations
+- Bug fixes are the best and always welcome! Like [this one](https://github.com/geohot/tinygrad/pull/421/files).
+- If you don't understand the code you are changing, don't change it!
+- All code golf PRs will be closed, but [conceptual cleanups](https://github.com/geohot/tinygrad/pull/372/files) are great.
+- Features are welcome. Though if you are adding a feature, you need to include tests.
+- Improving test coverage is great, with reliable non-brittle tests.
-```
-Buffer # class of memory on this device
-unary_op (NOOP, EXP, LOG, CAST, SIN) # A -> A
-reduce_op (SUM, MAX) # A -> B (smaller size, B has 1 in shape)
-binary_op (ADD, SUB, MUL, DIV, POW, CMPEQ, MAX) # A + A -> A (all the same size)
-movement_op (EXPAND, RESHAPE, PERMUTE, PAD, SHRINK, STRIDE) # A -> B (different size)
-fused_op [[optional]] (MULACC) # A * A -> B
-```
-
-## ImageNet inference
-
-Despite being tiny, tinygrad supports the full EfficientNet. Pass in a picture to discover what it is.
-
-```bash
-python3 examples/efficientnet.py https://media.istockphoto.com/photos/hen-picture-id831791190
-```
-
-Or, if you have a webcam and cv2 installed
-
-```bash
-python3 examples/efficientnet.py webcam
-```
-
-PROTIP: Set "DEBUG=2" environment variable if you want to see why it's slow.
-
-### tinygrad supports Stable Diffusion!
-
-You might need to download the [weight](https://huggingface.co/CompVis/stable-diffusion-v-1-4-original/resolve/main/sd-v1-4.ckpt) of Stable Diffusion and put it into weights/
-
-Run `python3 examples/stable_diffusion.py`
-
-
- 
-
-
-
-"a horse sized cat eating a bagel"
-
-
-### tinygrad supports LLaMA
-
-After putting the weights in weights/LLaMA, you can have a chat with Stacy. She lives inside tinygrad.
-
-```bash
-python3 examples/llama.py
-```
-
-### tinygrad supports GANs
-
-See `examples/mnist_gan.py`
-
-
- 
-
-
-### tinygrad supports yolo
-
-See `examples/yolov3.py`
-
-
- 
-
-
-### Drawing Execution Graph
-
-```bash
-GRAPH=1 python3 test/models/test_mnist.py TestMNIST.test_sgd_onestep
-# requires dot, outputs /tmp/net.svg
-```
+Additional guidelines can be found in [CONTRIBUTING.md](/CONTRIBUTING.md).
### Running tests
For more examples on how to run the full test suite please refer to the [CI workflow](.github/workflows/test.yml).
-```bash
+Some examples:
+```sh
python3 -m pip install -e '.[testing]'
python3 -m pytest
python3 -m pytest -v -k TestTrain
python3 ./test/models/test_train.py TestTrain.test_efficientnet
```
-
diff --git a/docs/README.md b/docs/README.md
index 052347e02e..6a8de5887c 100644
--- a/docs/README.md
+++ b/docs/README.md
@@ -1,125 +1,37 @@
-### Welcome to the tinygrad documentation
+# Welcome to the tinygrad documentation!
-General instructions you will find in [README.md](https://github.com/geohot/tinygrad/blob/master/README.md)
+Here you will find documentation for tinygrad, as well as some examples and tutorials.
-[abstraction.py](https://github.com/geohot/tinygrad/blob/master/docs/abstractions.py) is a well documented showcase of the abstraction stack.
+## Getting Started
-There are plenty of [tests](https://github.com/geohot/tinygrad/tree/master/test) you can read through
-[Examples](https://github.com/geohot/tinygrad/tree/master/examples) contains tinygrad implementations of popular models (vision and language) and neural networks. LLama, Stable diffusion, GANs and Yolo to name a few
+Read the quick start guide [here](/docs/quickstart.md).
-### Environment variables
-Here is a list of environment variables you can use with tinygrad.
-Most of these are self-explanatory, and used to enable an option at runtime.
-Example : `GPU=1 DEBUG=4 python3 -m pytest`
+Or if you want to jump right in to how tinygrad works, you can read the [abstraction stack](/docs/abstractions.py) documentation.
-The columns are: Variable, Value and Description
-They are also grouped into either general tinygrad or specific files
+Or if you want to see some examples, you can look at the examples in the [examples](/examples) directory.
-##### General tinygrad
-DEBUG: [1-4], enable debugging output, with 4 you get operations, timings, speed, generated code and more
-GPU: [1], enable the GPU backend
-CPU: [1], enable CPU backend
-MPS: [1], emable MPS device (for Mac M1 and after)
-METAL: [1], enable Metal backend (for Mac M1 and after)
-METAL_XCODE: [1], enable Metal using MacOS Xcode sdk
-TORCH: [1], enable Torch backend
-CLANG: [1], enable Clang backend
-LLVM: [1], enable LLVM backend
-LLVMOPT: [1], enable LLVM optimization
-LAZY: [1], enable lazy operations
-OPT: [1-4], enable optimization
-OPTLOCAL: [1], enable local optimization
-JIT: [1], enable Jit
-GRAPH: [1], Create a graph of all operations
-GRAPHPATH: [/path/to], what path to generate the graph image
-PRUNEGRAPH, [1], prune movementops and loadops from the graph
-PRINT_PRG: [1], print program
-FLOAT16: [1], use float16 instead of float32
-ENABLE_METHOD_CACHE: [1], enable method cache
-EARLY_STOPPING: [1], stop early
-DISALLOW_ASSIGN: [1], enable not assigning the realized lazydata to the lazy output buffer
+Or if you just want to see some of the things tinygrad can do, check out the [showcase](/docs/showcase.md).
-##### tinygrad/codegen/cstyle.py
-NATIVE_EXPLOG: [1], enable using native explog
+## API
-##### accel/ane/2_compile/hwx_parse.py
-PRINTALL: [1], print all ane registers
+This is currently a big work in progress.
-##### extra/onnx.py
-ONNXLIMIT: [ ], set a limit for Onnx
-DEBUGONNX: [1], enable Onnx debugging
+## Resources
-##### extra/thneed.py
-DEBUGCL: [1-4], enable Debugging for OpenCL
-PRINT_KERNEL: [1], Print OpenCL Kernels
+### Environment Variables
-##### extra/kernel_search.py
-OP: [1-3], different operations
-NOTEST: [1], enable not testing ast
-DUMP: [1], enable dumping of intervention cache
-REDUCE: [1], enable reduce operations
-SIMPLE_REDUCE: [1], enable simpler reduce operations
-BC: [1], enable big conv operations
-CONVW: [1], enable convw operations
-FASTCONV: [1], enable faster conv operations
-GEMM: [1], enable general matrix multiply operations
-BROKEN: [1], enable a kind of operation
-BROKEN3: [1], enable a kind of operation
+[env_vars.md](/docs/env_vars.md)
-##### examples/vit.py
-LARGE: [1], enable larger dimension model
+### Adding New Accelerators
-##### examples/llama.py
-WEIGHTS: [1], enable using weights
+[adding_new_accelerators.md](/docs/adding_new_accelerators.md)
-##### examples/mlperf
-MODEL: [resnet,retinanet,unet3d,rnnt,bert,maskrcnn], what models to use
+### Community
-##### examples/benchmark_train_efficientnet.py
-CNT: [10], the amount of times to loop the benchmark
-BACKWARD: [1], enable backward call
-TRAINING: [1], set Tensor.training
-CLCACHE: [1], enable Cache for OpenCL
+[](https://discord.gg/ZjZadyC7PK)
-##### examples/hlb_cifar10.py
-TORCHWEIGHTS: [1], use torch to initialize weights
-DISABLE_BACKWARD: [1], dont use backward operations
+## Contributing
-##### examples/benchmark_train_efficientnet.py & examples/hlb_cifar10.py
-ADAM: [1], enable Adam optimization
-
-##### examples/hlb_cifar10.py & xamples/hlb_cifar10_torch.py
-STEPS: [0-10], number of steps
-FAKEDATA: [1], enable to use random data
-
-##### examples/train_efficientnet.py
-STEPS: [1024 dividable], number of steps
-TINY: [1], use a tiny convolution network
-IMAGENET: [1], use imagenet for training
-
-##### examples/train_efficientnet.py & examples/train_resnet.py
-TRANSFER: [1], enable to use pretrained data
-
-##### examples & test/external/external_test_opt.py
-NUM: [18, 2], what ResNet[18] / EfficientNet[2] to train
-
-##### test/test_ops.py
-PRINT_TENSORS: [1], print tensors
-FORWARD_ONLY: [1], use forward operations only
-
-##### test/test_speed_v_torch.py
-TORCHCUDA: [1], enable the torch cuda backend
-
-##### test/external/external_test_gpu_ast.py
-KOPT: [1], enable kernel optimization
-KCACHE: [1], enable kernel cache
-
-##### test/external/external_test_opt.py
-ENET_NUM: [-2,-1], what EfficientNet to use
-
-##### test/test_dtype.py & test/extra/test_utils.py & extra/training.py
-CI: [1], enable to avoid some tests to run in CI
-
-##### examples & extra & test
-BS: [8, 16, 32, 64, 128], bytesize
+The documentation mainly follows the core contributing guidelines in the [README.md](/README.md#contributing).
+Additionally, we always welcome documentation contributions, especially for features that are currently under documented.
diff --git a/docs/abstractions.py b/docs/abstractions.py
index 06bb943e3b..28d9e514c9 100644
--- a/docs/abstractions.py
+++ b/docs/abstractions.py
@@ -98,16 +98,17 @@ class LazyOp:
src: Tuple[Union[LazyOp, LazyBuffer], ...] # the sources
arg: Optional[Any] = None # and an optional static argument
-# there's currently 20 Ops you have to implement for an accelerator.
-class UnaryOps(Enum): NOOP = auto(); EXP = auto(); LOG = auto(); NEG = auto(); NOT = auto()
-class BinaryOps(Enum): ADD = auto(); SUB = auto(); MUL = auto(); DIV = auto(); POW = auto(); CMPEQ = auto(); MAX = auto()
+# there's currently 27 Ops you have to implement for an accelerator.
+class UnaryOps(Enum): NOOP = auto(); EXP = auto(); LOG = auto(); CAST = auto(); SIN = auto()
+class BinaryOps(Enum): ADD = auto(); SUB = auto(); MUL = auto(); DIV = auto(); POW = auto(); CMPEQ = auto(); MAX = auto()
class ReduceOps(Enum): SUM = auto(); MAX = auto()
class MovementOps(Enum): RESHAPE = auto(); PERMUTE = auto(); EXPAND = auto(); PAD = auto(); SHRINK = auto(); STRIDE = auto()
-class LoadOps(Enum): FROMCPU = auto()
+class FusedOps(Enum): MULACC = auto()
+class LoadOps(Enum): EMPTY = auto(); RAND = auto(); CONST = auto(); FROM = auto(); CONTIGUOUS = auto(); CUSTOM = auto()
# NOTE: if you have a CompiledBuffer(DeviceBuffer)
# you do not need to implement the MovementOps
# as they are handled by the ShapeTracker(in tinygrad/shape/shapetracker.py, code 7/10)
-Op = Union[UnaryOps, BinaryOps, ReduceOps, MovementOps, LoadOps]
+Op = Union[UnaryOps, BinaryOps, ReduceOps, MovementOps, FusedOps, LoadOps]
# most of tinygrad/lazy.py is concerned with fusing Ops into LazyOps ASTs that map to GPUKernels
# it's beyond the scope of this tutorial, but you can read the file if interested
@@ -130,11 +131,11 @@ assert lazyop.op == BinaryOps.ADD
assert len(lazyop.src) == 2
# the first source is the 2, it comes from the CPU
-# the source is a LazyBuffer, since FROMCPU cannot be folded into LazyOp ASTs
+# the source is a LazyBuffer, holding the data as an ndarray
# again, a LazyOp AST is like a GPU kernel. you have to copy the data on the device first
print(lazyop.src[0].op)
-assert lazyop.src[0].op.op == LoadOps.FROMCPU
-assert lazyop.src[0].op.arg == [2], "the arg of the FROMCPU LazyOP is the [2.]"
+assert lazyop.src[0].op.op == LoadOps.FROM
+assert lazyop.src[0].op.src[0].realized.toCPU()[0] == 2, "the arg of the FROM LazyOP is a LazyBuffer holding [2.]"
assert result.lazydata.realized is None, "the LazyBuffer is not realized yet"
# now we realize the LazyBuffer
diff --git a/docs/adding_new_accelerators.md b/docs/adding_new_accelerators.md
new file mode 100644
index 0000000000..8957435cfb
--- /dev/null
+++ b/docs/adding_new_accelerators.md
@@ -0,0 +1,33 @@
+# Adding a new accelerator to tinygrad
+
+It's pretty easy to add a new accelerator to tinygrad. All you need to do is implement a total of 26 (optionally 27) low level ops. Then tinygrad takes care of the rest, handling derivatives and syntactic sugar.
+
+## llops
+
+These are the ops that you must implement for your accelerator of choice. Compiled Accelerators do not need to implement movement_ops, as they are handled b the ShapeTracker.
+```
+Buffer # class of memory on this device
+unary_op (NOOP, EXP, LOG, CAST, SIN) # A -> A
+reduce_op (SUM, MAX) # A -> B (smaller size, B has 1 in shape)
+binary_op (ADD, SUB, MUL, DIV, POW, CMPEQ, MAX) # A + A -> A (all the same size)
+movement_op (EXPAND, RESHAPE, PERMUTE, PAD, SHRINK, STRIDE) # A -> B (different size)
+load_op (EMPTY, RAND, CONST, FROM, CONTIGUOUS, CUSTOM) # -> A (initialize data on device)
+fused_op [[optional]] (MULACC) # A * A -> B
+```
+
+## mlops
+
+These are the mid level ops that handle the derivatives.
+```
+Relu, Log, Exp, Sin # unary ops
+Sum, Max # reduce ops (with axis argument)
+Maximum, Add, Sub, Mul, Pow, Div, Equal # binary ops (no broadcasting, use expand)
+Expand, Reshape, Permute, Pad, Shrink, Flip # movement ops
+```
+These are implemented in [mlops.py](/tinygrad/mlops.py).
+
+## hlops
+
+These are the syntax sugar. They are built on top of the mlops and support most of the things that you could expect from a tensor library.
+
+These are implemented in [tensor.py](/tinygrad/tensor.py).
diff --git a/docs/env_vars.md b/docs/env_vars.md
new file mode 100644
index 0000000000..e1cdbd27c5
--- /dev/null
+++ b/docs/env_vars.md
@@ -0,0 +1,186 @@
+# List of environment variables that control tinygrad behavior.
+
+This is a list of environment variable that control the runtime behavior of tinygrad and its examples.
+Most of these are self-explanatory, and are usually used to set an option at runtime.
+
+Example: `GPU=1 DEBUG=4 python3 -m pytest`
+
+The columns are: Variable, Possible Value(s) and Description.
+
+- A `#` means that the variable can take any integer value.
+
+## Global Variables
+
+These control the behavior of core tinygrad even when used as a library.
+
+Variable | Possible Value(s) | Description
+---|---|---
+DEBUG | [1-4] | enable debugging output, with 4 you get operations, timings, speed, generated code and more
+GPU | [1] | enable the GPU backend
+CPU | [1] | enable CPU backend
+MPS | [1] | enable MPS device (for Mac M1 and after)
+METAL | [1] | enable Metal backend (for Mac M1 and after)
+METAL_XCODE | [1] | enable Metal using macOS Xcode SDK
+TORCH | [1] | enable PyTorch backend
+CLANG | [1] | enable Clang backend
+LLVM | [1] | enable LLVM backend
+LLVMOPT | [1] | enable slightly more expensive LLVM optimizations
+LAZY | [1] | enable lazy operations (this is the default)
+OPT | [1-4] | optimization level
+OPTLOCAL | [1-2] | enable local optimization
+GRAPH | [1] | create a graph of all operations (requires graphviz)
+GRAPHPATH | [/path/to] | where to put the generated graph
+PRUNEGRAPH | [1] | prune MovementOps and LoadOps from the graph
+PRINT_PRG | [1] | print program code
+IMAGE | [1] | enable 2d specific optimizations
+FLOAT16 | [1] | use float16 for images instead of float32
+ENABLE_METHOD_CACHE | [1] | enable method cache (this is the default)
+EARLY_STOPPING | [# > 0] | stop after this many kernels
+DISALLOW_ASSIGN | [1] | disallow assignment of tensors
+NATIVE_EXPLOG | [1] | enable using native exp and log
+
+## File Specific Variables
+
+These are variables that control the behavior of a specific file, these usually don't affect the library itself.
+Most of the time these will never be used, but they are here for completeness.
+
+### accel/ane/2_compile/hwx_parse.py
+
+Variable | Possible Value(s) | Description
+---|---|---
+PRINTALL | [1] | print all ANE registers
+
+### extra/onnx.py
+
+Variable | Possible Value(s) | Description
+---|---|---
+ONNXLIMIT | [#] | set a limit for ONNX
+DEBUGONNX | [1] | enable ONNX debugging
+
+### extra/thneed.py
+
+Variable | Possible Value(s) | Description
+---|---|---
+DEBUGCL | [1-4] | enable Debugging for OpenCL
+PRINT_KERNEL | [1] | Print OpenCL Kernels
+
+### extra/kernel_search.py
+
+Variable | Possible Value(s) | Description
+---|---|---
+OP | [1-3] | different operations
+NOTEST | [1] | enable not testing AST
+DUMP | [1] | enable dumping of intervention cache
+REDUCE | [1] | enable reduce operations
+SIMPLE_REDUCE | [1] | enable simpler reduce operations
+BC | [1] | enable big conv operations
+CONVW | [1] | enable convw operations
+FASTCONV | [1] | enable faster conv operations
+GEMM | [1] | enable general matrix multiply operations
+BROKEN | [1] | enable a kind of operation
+BROKEN3 | [1] | enable a kind of operation
+
+### examples/vit.py
+
+Variable | Possible Value(s) | Description
+---|---|---
+LARGE | [1] | enable larger dimension model
+
+### examples/llama.py
+
+Variable | Possible Value(s) | Description
+---|---|---
+WEIGHTS | [1] | enable loading weights
+
+### examples/mlperf
+
+Variable | Possible Value(s) | Description
+---|---|---
+MODEL | [resnet,retinanet,unet3d,rnnt,bert,maskrcnn] | what models to use
+
+### examples/benchmark_train_efficientnet.py
+
+Variable | Possible Value(s) | Description
+---|---|---
+CNT | [10] | the amount of times to loop the benchmark
+BACKWARD | [1] | enable backward pass
+TRAINING | [1] | set Tensor.training
+CLCACHE | [1] | enable cache for OpenCL
+
+### examples/hlb_cifar10.py
+
+Variable | Possible Value(s) | Description
+---|---|---
+TORCHWEIGHTS | [1] | use torch to initialize weights
+DISABLE_BACKWARD | [1] | don't do backward pass
+
+### examples/benchmark_train_efficientnet.py & examples/hlb_cifar10.py
+
+Variable | Possible Value(s) | Description
+---|---|---
+ADAM | [1] | use the Adam optimizer
+
+### examples/hlb_cifar10.py & xamples/hlb_cifar10_torch.py
+
+Variable | Possible Value(s) | Description
+---|---|---
+STEPS | [0-10] | number of steps
+FAKEDATA | [1] | enable to use random data
+
+### examples/train_efficientnet.py
+
+Variable | Possible Value(s) | Description
+---|---|---
+STEPS | [# % 1024] | number of steps
+TINY | [1] | use a tiny convolution network
+IMAGENET | [1] | use imagenet for training
+
+### examples/train_efficientnet.py & examples/train_resnet.py
+
+Variable | Possible Value(s) | Description
+---|---|---
+TRANSFER | [1] | enable to use pretrained data
+
+### examples & test/external/external_test_opt.py
+
+Variable | Possible Value(s) | Description
+---|---|---
+NUM | [18, 2] | what ResNet[18] / EfficientNet[2] to train
+
+### test/test_ops.py
+
+Variable | Possible Value(s) | Description
+---|---|---
+PRINT_TENSORS | [1] | print tensors
+FORWARD_ONLY | [1] | use forward operations only
+
+### test/test_speed_v_torch.py
+
+Variable | Possible Value(s) | Description
+---|---|---
+TORCHCUDA | [1] | enable the torch cuda backend
+
+### test/external/external_test_gpu_ast.py
+
+Variable | Possible Value(s) | Description
+---|---|---
+KOPT | [1] | enable kernel optimization
+KCACHE | [1] | enable kernel cache
+
+### test/external/external_test_opt.py
+
+Variable | Possible Value(s) | Description
+---|---|---
+ENET_NUM | [-2,-1] | what EfficientNet to use
+
+### test/test_dtype.py & test/extra/test_utils.py & extra/training.py
+
+Variable | Possible Value(s) | Description
+---|---|---
+CI | [1] | disables some tests for CI
+
+### examples & extra & test
+
+Variable | Possible Value(s) | Description
+---|---|---
+BS | [8, 16, 32, 64, 128] | batch size to use
diff --git a/docs/quickstart.md b/docs/quickstart.md
new file mode 100644
index 0000000000..0b405f487e
--- /dev/null
+++ b/docs/quickstart.md
@@ -0,0 +1,300 @@
+# tinygrad Quick Start Guide
+
+This guide assumes no prior knowledge of pytorch or any other deep learning framework, but does assume some basic knowledge of neural networks.
+It is intended to be a very quick overview of the high level API that tinygrad provides.
+
+This guide is also structured as a tutorial which at the end of it you will have a working model that can classify handwritten digits.
+
+We need some imports to get started:
+```py
+import numpy as np
+import time
+```
+
+## Tensors
+
+Tensors are the base data structure in tinygrad. They can be thought of as a multidimensional array of a specific data type.
+All high level operations in tinygrad operate on these tensors.
+
+The tensor class can be imported like so:
+```py
+from tinygrad.tensor import Tensor
+```
+
+Tensors can be created from an existing data structure like a python list or numpy ndarray:
+```py
+t1 = Tensor([1, 2, 3, 4, 5])
+na = np.array([1, 2, 3, 4, 5])
+t2 = Tensor(na)
+```
+
+Tensors can also be created using one of the many factory methods:
+```py
+full = Tensor.full(shape=(2, 3), fill_value=5) # create a tensor of shape (2, 3) filled with 5
+zeros = Tensor.zeros(2, 3) # create a tensor of shape (2, 3) filled with 0
+ones = Tensor.ones(2, 3) # create a tensor of shape (2, 3) filled with 1
+
+full_like = Tensor.full_like(full, fill_value=2) # create a tensor of the same shape as `full` filled with 2
+zeros_like = Tensor.zeros_like(full) # create a tensor of the same shape as `full` filled with 0
+ones_like = Tensor.ones_like(full) # create a tensor of the same shape as `full` filled with 1
+
+eye = Tensor.eye(3) # create a 3x3 identity matrix
+arange = Tensor.arange(start=0, stop=10, step=1) # create a tensor of shape (10,) filled with values from 0 to 9
+
+rand = Tensor.rand(2, 3) # create a tensor of shape (2, 3) filled with random values from a uniform distribution
+randn = Tensor.randn(2, 3) # create a tensor of shape (2, 3) filled with random values from a normal distribution
+uniform = Tensor.uniform(2, 3, low=0, high=10) # create a tensor of shape (2, 3) filled with random values from a uniform distribution between 0 and 10
+```
+There are even more of these factory methods, you can find them in the [tensor.py](/tinygrad/tensor.py) file.
+
+All the tensors creation methods can take a `dtype` argument to specify the data type of the tensor.
+```py
+from tinygrad.helpers import dtypes
+
+t3 = Tensor([1, 2, 3, 4, 5], dtype=dtypes.int32)
+```
+
+Tensors allow you to perform operations on them like so:
+```py
+t4 = Tensor([1, 2, 3, 4, 5])
+t5 = (t4 + 1) * 2
+t6 = (t5 * t4).relu().log_softmax()
+```
+
+All of these operations are lazy and are only executed when you realize the tensor using `.realize()` or `.numpy()`.
+```py
+print(t6.numpy())
+# [-56. -48. -36. -20. 0.]
+```
+
+There are a lot more operations that can be performed on tensors, you can find them in the [tensor.py](/tinygrad/tensor.py) file.
+Additionally reading through [abstractions.py](/docs/abstractions.py) will help you understand how operations on these tensors make their way down to your hardware.
+
+## Models
+
+Neural networks in tinygrad are really just represented by the operations performed on tensors.
+These operations are commonly grouped into the `__call__` method of a class which allows modularization and reuse of these groups of operations.
+These classes do not need to inherit from any base class, in fact if they don't need any trainable parameters they don't even need to be a class!
+
+An example of this would be the `nn.Linear` class which represents a linear layer in a neural network.
+```py
+# from tinygrad.nn import Linear
+class Linear:
+ def __init__(self, in_features, out_features, bias=True, initialization: str='kaiming_uniform'):
+ self.weight = getattr(Tensor, initialization)(out_features, in_features)
+ self.bias = Tensor.zeros(out_features) if bias else None
+
+ def __call__(self, x):
+ return x.linear(self.weight.transpose(), self.bias)
+```
+There are more neural network modules already implemented in [nn](/tinygrad/nn/__init__.py), and you can also implement your own.
+
+We will be implementing a simple neural network that can classify handwritten digits from the MNIST dataset.
+Our classifier will be a simple 2 layer neural network with a Leaky ReLU activation function.
+It will use a hidden layer size of 128 and an output layer size of 10 (one for each digit) with no bias on either Linear layer.
+```py
+from tinygrad.nn import Linear
+
+class TinyNet:
+ def __init__(self):
+ self.l1 = Linear(784, 128, bias=False)
+ self.l2 = Linear(128, 10, bias=False)
+
+ def __call__(self, x):
+ x = self.l1(x)
+ x = x.leakyrelu()
+ x = self.l2(x)
+ return x.log_softmax()
+
+net = TinyNet()
+```
+We can see that the forward pass of our neural network is just the sequence of operations performed on the input tensor `x`.
+We can also see that functional operations like `leakyrelu` and `log_softmax` are not defined as classes and instead are just methods we can just call.
+Finally, we just initialize an instance of our neural network, and we are ready to start training it.
+
+## Training
+
+Now that we have our neural network defined we can start training it.
+Training neural networks in tinygrad is super simple.
+All we need to do is define our neural network, define our loss function, and then call `.backward()` on the loss function to compute the gradients.
+They can then be used to update the parameters of our neural network using one of the many optimizers in [optim.py](/tinygrad/nn/optim.py).
+
+First we need to set the training flag in `Tensor`:
+```py
+Tensor.training = True
+```
+
+For our loss function we will be using cross entropy loss.
+```py
+# from extra.training import sparse_categorical_crossentropy
+def cross_entropy(out, Y):
+ num_classes = out.shape[-1]
+ YY = Y.flatten().astype(np.int32)
+ y = np.zeros((YY.shape[0], num_classes), np.float32)
+ y[range(y.shape[0]),YY] = -1.0*num_classes
+ y = y.reshape(list(Y.shape)+[num_classes])
+ y = Tensor(y)
+ return out.mul(y).mean()
+```
+As we can see in this implementation of cross entropy loss, there are certain operations that tinygrad does not support.
+Namely, operations that are load/store like indexing a tensor with another tensor or assigning a value to a tensor at a certain index.
+Load/store ops are not supported in tinygrad because they add complexity when trying to port to different backends and 90% of the models out there don't use/need them.
+
+For our optimizer we will be using the traditional stochastic gradient descent optimizer with a learning rate of 3e-4.
+```py
+from tinygrad.nn.optim import SGD
+
+opt = SGD([net.l1.weight, net.l2.weight], lr=3e-4)
+```
+We can see that we are passing in the parameters of our neural network to the optimizer.
+This is due to the fact that the optimizer needs to know which parameters to update.
+There is a simpler way to do this just by using `get_parameters(net)` from `tinygrad.nn.optim` which will return a list of all the parameters in the neural network.
+The parameters are just listed out explicitly here for clarity.
+
+Now that we have our network, loss function, and optimizer defined all we are missing is the data to train on!
+There are a couple of dataset loaders in tinygrad located in [/datasets](/datasets).
+We will be using the MNIST dataset loader.
+```py
+from datasets import fetch_mnist
+```
+
+Now we have everything we need to start training our neural network.
+We will be training for 1000 steps with a batch size of 64.
+```py
+X_train, Y_train, X_test, Y_test = fetch_mnist()
+
+for step in range(1000):
+ # random sample a batch
+ samp = np.random.randint(0, X_train.shape[0], size=(64))
+ batch = Tensor(X_train[samp], requires_grad=False)
+ # get the corresponding labels
+ labels = Y_train[samp]
+
+ # forward pass
+ out = net(batch)
+
+ # compute loss
+ loss = cross_entropy(out, labels)
+
+ # zero gradients
+ opt.zero_grad()
+
+ # backward pass
+ loss.backward()
+
+ # update parameters
+ opt.step()
+
+ # calculate accuracy
+ pred = np.argmax(out.numpy(), axis=-1)
+ acc = (pred == labels).mean()
+
+ if step % 100 == 0:
+ print(f"Step {step+1} | Loss: {loss.numpy()} | Accuracy: {acc}")
+```
+
+## Evaluation
+
+Now that we have trained our neural network we can evaluate it on the test set.
+We will be using the same batch size of 64 and will be evaluating for 1000 of those batches.
+```py
+# set training flag to false
+Tensor.training = False
+
+st = time.perf_counter()
+avg_acc = 0
+for step in range(1000):
+ # random sample a batch
+ samp = np.random.randint(0, X_test.shape[0], size=(64))
+ batch = Tensor(X_test[samp], requires_grad=False)
+ # get the corresponding labels
+ labels = Y_test[samp]
+
+ # forward pass
+ out = net(batch)
+
+ # calculate accuracy
+ pred = np.argmax(out.numpy(), axis=-1)
+ avg_acc += (pred == labels).mean()
+print(f"Test Accuracy: {avg_acc / 1000}")
+print(f"Time: {time.perf_counter() - st}")
+```
+
+## And that's it!
+
+Highly recommend you check out the [examples/](/examples) folder for more examples of using tinygrad.
+Reading the source code of tinygrad is also a great way to learn how it works.
+Specifically the tests in [tests/](/tests) are a great place to see how to use and the semantics of the different operations.
+There are also a bunch of models implemented in [models/](/models) that you can use as a reference.
+
+Additionally, feel free to ask questions in the `#learn-tinygrad` channel on the [discord](https://discord.gg/beYbxwxVdx). Don't ask to ask, just ask!
+
+## Extras
+
+### JIT
+
+Additionally, it is possible to speed up the computation of certain neural networks by using the JIT.
+Currently, this does not support models with varying input sizes and non tinygrad operations.
+
+To use the JIT we just need to add a function decorator to the forward pass of our neural network and ensure that the input and output are realized tensors.
+Or in this case we will create a wrapper function and decorate the wrapper function to speed up the evaluation of our neural network.
+```py
+from tinygrad.jit import TinyJit
+
+@TinyJit
+def jit(x):
+ return net(x).realize()
+
+st = time.perf_counter()
+avg_acc = 0
+for step in range(1000):
+ # random sample a batch
+ samp = np.random.randint(0, X_test.shape[0], size=(64))
+ batch = Tensor(X_test[samp], requires_grad=False)
+ # get the corresponding labels
+ labels = Y_test[samp]
+
+ # forward pass with jit
+ out = jit(batch)
+
+ # calculate accuracy
+ pred = np.argmax(out.numpy(), axis=-1)
+ avg_acc += (pred == labels).mean()
+print(f"Test Accuracy: {avg_acc / 1000}")
+print(f"Time: {time.perf_counter() - st}")
+```
+You will find that the evaluation time is much faster than before and that your accelerator utilization is much higher.
+
+### Saving and Loading Models
+
+The standard weight format for tinygrad is [safetensors](https://github.com/huggingface/safetensors). This means that you can load the weights of any model also using safetensors into tinygrad.
+There are functions in [state.py](/tinygrad/state.py) to save and load models to and from this format.
+```py
+from tinygrad.state import safe_save, safe_load, get_state_dict, load_state_dict
+
+# first we need the state dict of our model
+state_dict = get_state_dict(net)
+
+# then we can just save it to a file
+safe_save(state_dict, "model.safetensors")
+
+# and load it back in
+state_dict = safe_load("model.safetensors")
+load_state_dict(net, state_dict)
+```
+
+Many of the models in the [models/](/models) folder have a `load_from_pretrained` method that will download and load the weights for you. These usually are pytorch weights meaning that you would need pytorch installed to load them.
+
+### Environment Variables
+
+There exist a bunch of environment variables that control the runtime behavior of tinygrad.
+Some of the commons ones are `DEBUG` and the different backend enablement variables.
+
+You can find a full list and their descriptions in [env_vars.md](/docs/env_vars.md).
+
+### Visualizing the Computation Graph
+
+It is possible to visualize the computation graph of a neural network using [graphviz](https://graphviz.org/).
+
+This is easily done by running a single pass (forward or backward!) of the neural network with the environment variable `GRAPH` set to `1`.
+The graph will be saved to `/tmp/net.svg` by default.
diff --git a/docs/showcase.md b/docs/showcase.md
new file mode 100644
index 0000000000..38bc9b9594
--- /dev/null
+++ b/docs/showcase.md
@@ -0,0 +1,59 @@
+# tinygrad Showcase
+
+Despite being a tiny library, tinygrad is capable of doing a lot of things. From state-of-the-art [vision](https://arxiv.org/abs/1905.11946) to state-of-the-art [language](https://arxiv.org/abs/1706.03762) models.
+
+## Vision
+
+### EfficientNet
+
+You can either pass in the URL of a picture to discover what it is:
+```sh
+python3 examples/efficientnet.py https://media.istockphoto.com/photos/hen-picture-id831791190
+```
+Or, if you have a camera and OpenCV installed, you can detect what is in front of you:
+```sh
+python3 examples/efficientnet.py webcam
+```
+
+### YOLOv3
+
+Take a look at [yolov3.py](/examples/yolov3.py).
+
+
+
+## Audio
+
+### Whisper
+
+Take a look at [whisper.py](/examples/whisper.py). You need pyaudio and torchaudio installed.
+
+```sh
+SMALL=1 python3 examples/whisper.py
+```
+
+## Generative
+
+### Generative Adversarial Networks
+
+Take a look at [mnist_gan.py](/examples/mnist_gan.py).
+
+
+
+### Stable Diffusion
+
+```sh
+python3 examples/stable_diffusion.py
+```
+
+
+
+*"a horse sized cat eating a bagel"*
+
+### LLaMA
+
+You will need to download and put the weights into the `weights/LLaMA` directory, which may need to be created.
+
+Then you can have a chat with Stacy:
+```sh
+python3 examples/llama.py
+```
diff --git a/docs/mnist_by_tinygrad.jpg b/docs/showcase/mnist_by_tinygrad.jpg
similarity index 100%
rename from docs/mnist_by_tinygrad.jpg
rename to docs/showcase/mnist_by_tinygrad.jpg
diff --git a/docs/stable_diffusion_by_tinygrad.jpg b/docs/showcase/stable_diffusion_by_tinygrad.jpg
similarity index 100%
rename from docs/stable_diffusion_by_tinygrad.jpg
rename to docs/showcase/stable_diffusion_by_tinygrad.jpg
diff --git a/docs/yolo_by_tinygrad.jpg b/docs/showcase/yolo_by_tinygrad.jpg
similarity index 100%
rename from docs/yolo_by_tinygrad.jpg
rename to docs/showcase/yolo_by_tinygrad.jpg
diff --git a/examples/whisper.py b/examples/whisper.py
new file mode 100644
index 0000000000..40e4009c06
--- /dev/null
+++ b/examples/whisper.py
@@ -0,0 +1,247 @@
+# thanks to https://github.com/openai/whisper for a good chunk of MIT licensed code
+
+import sys
+import pathlib
+import base64
+import multiprocessing
+import numpy as np
+from typing import Optional
+from extra.utils import download_file
+from tinygrad.state import torch_load, load_state_dict
+from tinygrad.helpers import getenv
+import tinygrad.nn as nn
+from tinygrad.tensor import Tensor
+
+# TODO: you have written this fifteen times
+class MultiHeadAttention:
+ def __init__(self, n_state, n_head):
+ self.n_head = n_head
+ self.query = nn.Linear(n_state, n_state)
+ self.key = nn.Linear(n_state, n_state, bias=False)
+ self.value = nn.Linear(n_state, n_state)
+ self.out = nn.Linear(n_state, n_state)
+
+ def __call__(self, x:Tensor, xa:Optional[Tensor]=None, mask:Optional[Tensor]=None):
+ q = self.query(x)
+ k = self.key(xa or x)
+ v = self.value(xa or x)
+ wv, qk = self.qkv_attention(q, k, v, mask)
+ # NOTE: we aren't returning qk
+ return self.out(wv)
+
+ def qkv_attention(self, q, k, v, mask=None):
+ n_batch, n_ctx, n_state = q.shape
+ scale = (n_state // self.n_head) ** -0.25
+ q = q.reshape(*q.shape[:2], self.n_head, -1).permute(0, 2, 1, 3) * scale
+ k = k.reshape(*k.shape[:2], self.n_head, -1).permute(0, 2, 3, 1) * scale
+ v = v.reshape(*v.shape[:2], self.n_head, -1).permute(0, 2, 1, 3)
+ qk = q @ k
+ if mask is not None: qk = qk + mask[:n_ctx, :n_ctx]
+ w = qk.softmax(-1)
+ return (w @ v).permute(0, 2, 1, 3).flatten(start_dim=2), qk.detach()
+
+class ResidualAttentionBlock:
+ def __init__(self, n_state, n_head, cross_attention=False):
+ self.attn = MultiHeadAttention(n_state, n_head)
+ self.attn_ln = nn.LayerNorm(n_state)
+
+ self.cross_attn = MultiHeadAttention(n_state, n_head) if cross_attention else None
+ self.cross_attn_ln = nn.LayerNorm(n_state) if cross_attention else None
+
+ self.mlp = [nn.Linear(n_state, n_state*4), Tensor.gelu, nn.Linear(n_state*4, n_state)]
+ self.mlp_ln = nn.LayerNorm(n_state)
+
+ def __call__(self, x, xa=None, mask=None):
+ x = x + self.attn(self.attn_ln(x), mask=mask)
+ if self.cross_attn: x = x + self.cross_attn(self.cross_attn_ln(x), xa)
+ x = x + self.mlp_ln(x).sequential(self.mlp)
+ return x
+
+class AudioEncoder:
+ def __init__(self, n_mels, n_audio_ctx, n_audio_state, n_audio_head, n_audio_layer, **_):
+ self.conv1 = nn.Conv1d(n_mels, n_audio_state, kernel_size=3, padding=1)
+ self.conv2 = nn.Conv1d(n_audio_state, n_audio_state, kernel_size=3, stride=2, padding=1)
+ self.blocks = [ResidualAttentionBlock(n_audio_state, n_audio_head) for _ in range(n_audio_layer)]
+ self.ln_post = nn.LayerNorm(n_audio_state)
+ self.positional_embedding = Tensor.empty(n_audio_ctx, n_audio_state)
+
+ def __call__(self, x):
+ x = self.conv1(x).gelu()
+ x = self.conv2(x).gelu()
+ x = x.permute(0, 2, 1)
+ x = x + self.positional_embedding[:x.shape[1]]
+ x = x.sequential(self.blocks)
+ x = self.ln_post(x)
+ return x
+
+class TextDecoder:
+ def __init__(self, n_vocab, n_text_ctx, n_text_state, n_text_head, n_text_layer, **_):
+ self.token_embedding = nn.Embedding(n_vocab, n_text_state)
+ self.positional_embedding = Tensor.empty(n_text_ctx, n_text_state)
+ self.blocks = [ResidualAttentionBlock(n_text_state, n_text_head, cross_attention=True) for _ in range(n_text_layer)]
+ self.ln = nn.LayerNorm(n_text_state)
+ #mask = torch.empty(n_ctx, n_ctx).fill_(-np.inf).triu_(1)
+
+ def __call__(self, x, xa):
+ offset = 0
+ x = self.token_embedding(x) + self.positional_embedding[offset : offset + x.shape[-1]]
+
+ seqlen, start_pos = x.shape[1], 0
+
+ mask = np.full((1, 1, seqlen, start_pos + seqlen), float("-inf"), dtype=np.float32)
+ mask = np.triu(mask, k=start_pos + 1) # TODO: this is hard to do in tinygrad
+ mask = Tensor(mask)
+
+ for block in self.blocks: x = block(x, xa, mask)
+ x = self.ln(x)
+ return x @ self.token_embedding.weight.T
+
+class Whisper:
+ def __init__(self, dims):
+ self.encoder = AudioEncoder(**dims)
+ self.decoder = TextDecoder(**dims)
+
+ def __call__(self, mel:Tensor, tokens:Tensor):
+ return self.decoder(tokens, self.encoder(mel))
+
+# TODO: this is tragic. remove this
+import functools
+import torch
+import torchaudio
+import librosa
+
+@functools.lru_cache(None)
+def get_filters(sample_rate, n_fft, n_mels):return torch.tensor(librosa.filters.mel(sr=sample_rate, n_fft=n_fft, n_mels=n_mels))
+@functools.lru_cache(None)
+def get_window(n_fft): return torch.hann_window(n_fft)
+
+def prep_audio(waveform, sample_rate) -> Tensor:
+ N_FFT = 400
+ HOP_LENGTH = 160
+ N_MELS = 80
+ stft = torch.stft(waveform, N_FFT, HOP_LENGTH, window=get_window(N_FFT), return_complex=True)
+ magnitudes = stft[..., :-1].abs() ** 2
+ mel_spec = get_filters(sample_rate, N_FFT, N_MELS) @ magnitudes
+ log_spec = torch.clamp(mel_spec, min=1e-10).log10()
+ log_spec = torch.maximum(log_spec, log_spec.max() - 8.0)
+ log_spec = (log_spec + 4.0) / 4.0
+ #print(waveform.shape, log_spec.shape)
+ return log_spec.numpy()
+
+LANGUAGES = {
+ "en": "english", "zh": "chinese", "de": "german", "es": "spanish", "ru": "russian", "ko": "korean", "fr": "french", "ja": "japanese", "pt": "portuguese", "tr": "turkish",
+ "pl": "polish", "ca": "catalan", "nl": "dutch", "ar": "arabic", "sv": "swedish", "it": "italian", "id": "indonesian", "hi": "hindi", "fi": "finnish", "vi": "vietnamese",
+ "he": "hebrew", "uk": "ukrainian", "el": "greek", "ms": "malay", "cs": "czech", "ro": "romanian", "da": "danish", "hu": "hungarian", "ta": "tamil", "no": "norwegian",
+ "th": "thai", "ur": "urdu", "hr": "croatian", "bg": "bulgarian", "lt": "lithuanian", "la": "latin", "mi": "maori", "ml": "malayalam", "cy": "welsh", "sk": "slovak", "te": "telugu",
+ "fa": "persian", "lv": "latvian", "bn": "bengali", "sr": "serbian", "az": "azerbaijani", "sl": "slovenian", "kn": "kannada", "et": "estonian", "mk": "macedonian",
+ "br": "breton", "eu": "basque", "is": "icelandic", "hy": "armenian", "ne": "nepali", "mn": "mongolian", "bs": "bosnian", "kk": "kazakh", "sq": "albanian", "sw": "swahili",
+ "gl": "galician", "mr": "marathi", "pa": "punjabi", "si": "sinhala", "km": "khmer", "sn": "shona", "yo": "yoruba", "so": "somali", "af": "afrikaans", "oc": "occitan", "ka": "georgian",
+ "be": "belarusian", "tg": "tajik", "sd": "sindhi", "gu": "gujarati", "am": "amharic", "yi": "yiddish", "lo": "lao", "uz": "uzbek", "fo": "faroese", "ht": "haitian creole",
+ "ps": "pashto", "tk": "turkmen", "nn": "nynorsk", "mt": "maltese", "sa": "sanskrit", "lb": "luxembourgish", "my": "myanmar", "bo": "tibetan", "tl": "tagalog", "mg": "malagasy",
+ "as": "assamese", "tt": "tatar", "haw": "hawaiian", "ln": "lingala", "ha": "hausa", "ba": "bashkir", "jw": "javanese", "su": "sundanese",
+}
+
+BASE = pathlib.Path(__file__).parent.parent / "weights"
+def get_encoding(n_vocab_in):
+ download_file("https://raw.githubusercontent.com/openai/whisper/main/whisper/assets/gpt2.tiktoken", BASE / "gpt2.tiktoken")
+ ranks = {base64.b64decode(token): int(rank) for token, rank in (line.split() for line in open(BASE / "gpt2.tiktoken") if line)}
+ n_vocab = len(ranks)
+ specials = [
+ "<|endoftext|>",
+ "<|startoftranscript|>",
+ *[f"<|{lang}|>" for lang in LANGUAGES.keys()],
+ "<|translate|>",
+ "<|transcribe|>",
+ "<|startoflm|>",
+ "<|startofprev|>",
+ "<|nospeech|>",
+ "<|notimestamps|>",
+ *[f"<|{i * 0.02:.2f}|>" for i in range(1501)],
+ ]
+ special_tokens = {}
+ for token in specials:
+ special_tokens[token] = n_vocab
+ n_vocab += 1
+ assert n_vocab == n_vocab_in
+ import tiktoken
+ return tiktoken.Encoding(
+ name="bob",
+ explicit_n_vocab=n_vocab,
+ pat_str=r"""'s|'t|'re|'ve|'m|'ll|'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+""",
+ mergeable_ranks=ranks,
+ special_tokens=special_tokens)
+
+def img(x):
+ import matplotlib.pyplot as plt
+ plt.imshow(x.numpy())
+ plt.show()
+
+RATE = 16000
+CHUNK = 1600
+RECORD_SECONDS = 10
+
+def listener(q):
+ prep_audio(torch.zeros(300), RATE)
+ import pyaudio
+ p = pyaudio.PyAudio()
+ stream = p.open(format=pyaudio.paInt16, channels=1, rate=RATE, input=True, frames_per_buffer=CHUNK)
+ print("listening")
+ for _ in range(0, int(RATE / CHUNK * RECORD_SECONDS)):
+ data = stream.read(CHUNK)
+ waveform = ((np.frombuffer(data, np.int16)/32768).astype(np.float32)*3).reshape(1, -1)
+ q.put(waveform)
+ print("done listening")
+
+if __name__ == "__main__":
+ if getenv("SMALL"):
+ fn = BASE / "whisper-small.en.pt"
+ download_file("https://openaipublic.azureedge.net/main/whisper/models/f953ad0fd29cacd07d5a9eda5624af0f6bcf2258be67c92b79389873d91e0872/small.en.pt", fn)
+ else:
+ fn = BASE / "whisper-tiny.en.pt"
+ download_file("https://openaipublic.azureedge.net/main/whisper/models/d3dd57d32accea0b295c96e26691aa14d8822fac7d9d27d5dc00b4ca2826dd03/tiny.en.pt", fn)
+ state = torch_load(fn)
+ model = Whisper(state['dims'])
+ load_state_dict(model, state['model_state_dict'])
+ enc = get_encoding(state['dims']['n_vocab'])
+
+ if len(sys.argv) > 1:
+ # offline
+ waveform, sample_rate = torchaudio.load(sys.argv[1], normalize=True)
+ log_spec = prep_audio(waveform, sample_rate)
+ lst = [enc._special_tokens["<|startoftranscript|>"]]
+ dat = model.encoder(Tensor(log_spec)).realize()
+ for i in range(50):
+ out = model.decoder(Tensor([lst]), dat)
+ out.realize()
+ idx = out[0,-1].numpy().argmax()
+ lst.append(idx)
+ print(enc.decode(lst))
+ else:
+ # online
+
+ q = multiprocessing.Queue()
+ p = multiprocessing.Process(target=listener, args=(q,))
+ p.daemon = True
+ p.start()
+
+ lst = [enc._special_tokens["<|startoftranscript|>"]]
+ total = None
+ did_read = False
+ for i in range(0, int(RATE / CHUNK * RECORD_SECONDS)):
+ while not q.empty() or total is None:
+ waveform = q.get()
+ if total is None: total = waveform
+ else: total = np.concatenate([total, waveform], axis=1)
+ did_read = True
+ if did_read:
+ last_total = total.shape[1]
+ log_spec = prep_audio(torch.Tensor(total), RATE)
+ encoded_audio = model.encoder(Tensor(log_spec)).realize()
+ out = model.decoder(Tensor([lst]), encoded_audio).realize()
+ idx = out[0,-1].numpy().argmax()
+ lst.append(idx)
+ dec = enc.decode(lst)
+ print(dec) # DO NOT REMOVE PRINT. IT'S VERY IMPORTANT
+ if dec.endswith("<|endoftext|>"):
+ #total = total[:, 320*(len(lst)-1):]
+ lst = [enc._special_tokens["<|startoftranscript|>"]]
diff --git a/extra/intel/.gitignore b/extra/intel/.gitignore
new file mode 100644
index 0000000000..cba7efc8ef
--- /dev/null
+++ b/extra/intel/.gitignore
@@ -0,0 +1 @@
+a.out
diff --git a/extra/intel/README b/extra/intel/README
new file mode 100644
index 0000000000..6a6ed7cca2
--- /dev/null
+++ b/extra/intel/README
@@ -0,0 +1,2 @@
+source /opt/intel/oneapi/compiler/latest/env/vars.sh
+sycl-ls
diff --git a/extra/intel/benchmark_matmul.py b/extra/intel/benchmark_matmul.py
new file mode 100644
index 0000000000..5999039de0
--- /dev/null
+++ b/extra/intel/benchmark_matmul.py
@@ -0,0 +1,57 @@
+import time
+
+onnx_path = "/tmp/my.onnx"
+N = 2048
+CNT = 400
+
+"""
+import torch
+import torch.nn as nn
+#dtype = torch.bfloat16
+dtype = torch.float32
+class MatMul(nn.Module):
+ def __init__(self):
+ super().__init__()
+ self.a = nn.Linear(N, N, bias=False)
+ def forward(self, x):
+ x = x.to(dtype)
+ for i in range(CNT): x = self.a(x).relu()
+ return x.to(torch.float32)
+
+torch_model = MatMul().to(dtype)
+torch.onnx.export(torch_model, torch.randn(N, N), onnx_path)
+"""
+
+"""
+import onnx
+from tinygrad.tensor import Tensor
+from extra.onnx import get_run_onnx
+out = get_run_onnx(onnx.load(onnx_path))({"onnx::MatMul_0": Tensor.zeros(N, N)})
+for x in out.values(): x.realize()
+"""
+
+from openvino.runtime import Core
+core = Core()
+devices = core.available_devices
+for device in devices:
+ device_name = core.get_property(device, "FULL_DEVICE_NAME")
+ print(f"{device}: {device_name}")
+model = core.read_model(onnx_path)
+compiled_model = core.compile_model(model, device_name='GPU.0')
+print(compiled_model)
+ireq = compiled_model.create_infer_request()
+for model_input in compiled_model.inputs:
+ tensor = ireq.get_tensor(model_input)
+ tensor.data[:] = 2
+ print(tensor)
+print("request")
+ireq.infer()
+ireq.infer()
+print("did one")
+
+REPS = 20
+st = time.perf_counter()
+for i in range(REPS): ireq.infer()
+et = time.perf_counter() - st
+print(f"{et*1000:.2f} ms {(CNT*N*N*N*REPS*2/et)*1e-9:.2f} GFLOPS")
+
diff --git a/extra/intel/go.sh b/extra/intel/go.sh
new file mode 100755
index 0000000000..8c67088c05
--- /dev/null
+++ b/extra/intel/go.sh
@@ -0,0 +1,3 @@
+#!/bin/bash -e
+/opt/intel/oneapi/compiler/latest/linux/bin-llvm/clang++ joint_matrix_bfloat16.cpp -fsycl
+SYCL_PI_TRACE=1 ./a.out
diff --git a/extra/intel/joint_matrix_bfloat16.cpp b/extra/intel/joint_matrix_bfloat16.cpp
new file mode 100644
index 0000000000..b21d6089d2
--- /dev/null
+++ b/extra/intel/joint_matrix_bfloat16.cpp
@@ -0,0 +1,173 @@
+//==-------- joint_matrix_bfloat16.cpp - DPC++ joint_matrix----------- ----==//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+// REQUIRES: matrix
+
+// RUN: %clangxx -fsycl %s -o %t.out -DSYCL_EXT_ONEAPI_MATRIX_VERSION=4
+// RUN: %CPU_RUN_PLACEHOLDER %t.out
+// RUN: %GPU_RUN_PLACEHOLDER %t.out
+
+#include
+#include
+
+using namespace sycl;
+using namespace sycl::ext::oneapi::experimental::matrix;
+using bfloat16 = sycl::ext::oneapi::bfloat16;
+
+//#define SG_SZ 16
+#define SG_SZ 8
+
+#define TM 8
+#define TN SG_SZ
+//#define TK 16
+#define TK 16
+
+#define BF16_EPSILON 0.00781250
+
+template struct big_matrix {
+private:
+ T *mat;
+
+public:
+ T *get_data() { return mat; }
+ void set_data(T *data) { mat = data; }
+ big_matrix(T *data) : mat(data) {}
+};
+
+template
+void matrix_multiply(big_matrix &C, big_matrix &A, big_matrix &B) {
+ size_t NDRangeM = M / TM;
+ size_t NDRangeN = N / TN;
+ buffer bufA(A.get_data(), range<2>(M, K));
+ buffer bufB(B.get_data(), range<2>(K, N));
+ buffer bufC((float *)C.get_data(), range<2>(M, N));
+
+ auto program = [&](handler &cgh) {
+ auto accC = bufC.get_access(cgh);
+ auto accA = bufA.get_access(cgh);
+ auto accB = bufB.get_access(cgh);
+
+ cgh.parallel_for(
+ nd_range<2>({NDRangeM, NDRangeN * SG_SZ}, {1, 1 * SG_SZ}),
+ [=](nd_item<2> spmd_item) [[intel::reqd_sub_group_size(SG_SZ)]]
+ {
+ // The submatrix API has to be accessed by all the workitems in a
+ // subgroup these functions will be called once by the subgroup no
+ // code divergence between the workitems
+ const auto global_idx = spmd_item.get_global_id(0);
+ const auto global_idy = spmd_item.get_global_id(1);
+ const auto sg_startx = global_idx - spmd_item.get_local_id(0);
+ const auto sg_starty = global_idy - spmd_item.get_local_id(1);
+
+ sub_group sg = spmd_item.get_sub_group();
+ joint_matrix sub_a;
+ // For B, we assume B has been already VNNIed.
+ joint_matrix sub_b;
+ joint_matrix sub_c;
+ joint_matrix_load(sg, sub_c, accC.get_pointer() + (sg_startx * TM) * N + sg_starty / SG_SZ * TN, N, layout::row_major);
+
+ for (int k = 0; k < K / TK; k += 1) { //
+ joint_matrix_load(sg, sub_a, accA.get_pointer() + (sg_startx * TM) * K + k * TK, K);
+ joint_matrix_load(sg, sub_b, accB.get_pointer() + (k * TK / 2) * (N * 2) + sg_starty / SG_SZ * TN * 2, N * 2);
+ sub_c = joint_matrix_mad(sg, sub_a, sub_b, sub_c);
+ }
+ joint_matrix_store(sg, sub_c, accC.get_pointer() + (sg_startx * TM) * N + sg_starty / SG_SZ * TN, N, layout::row_major);
+ }); // parallel for
+ };
+
+ queue q;
+ auto start = std::chrono::steady_clock::now();
+ auto e = q.submit(program);
+ auto submit = std::chrono::steady_clock::now();
+ e.wait();
+ auto end = std::chrono::steady_clock::now();
+ std::cout << "submit: " << std::chrono::duration_cast(submit - start).count() << " ms" << std::endl;
+ std::cout << "compute: " << std::chrono::duration_cast(end - submit).count() << " ms" << std::endl;
+
+ // ahh, freeing is slow
+}
+
+//#define SCALE 1024
+//#define SCALE 64
+#define SCALE 256
+static constexpr size_t MATRIX_M = TM * SCALE;
+static constexpr size_t MATRIX_N = TN * SCALE;
+static constexpr size_t MATRIX_K = TK * SCALE;
+bfloat16 A[MATRIX_M][MATRIX_K];
+bfloat16 B[MATRIX_K / 2][MATRIX_N * 2];
+float C[MATRIX_M][MATRIX_N];
+float D[MATRIX_M][MATRIX_N];
+
+float make_fp32(bfloat16 x) {
+ unsigned int y = *((int *)&x);
+ y = y << 16;
+ float *res = reinterpret_cast(&y);
+ return *res;
+}
+
+void matrix_multiply_ref(int *A_mem, int *B_mem, int *C_mem, int M, int N,
+ int K) {
+ for (int m = 0; m < M; m++)
+ for (int n = 0; n < N; n++) {
+ for (int k = 0; k < K; k++) {
+ // Because B was assumed VNNIed
+ bfloat16 *va = (bfloat16 *)(A_mem + m * K + k);
+ bfloat16 *vb = (bfloat16 *)(B_mem + k * N + n);
+ float acc = *((float *)(C_mem + m * N + n));
+ for (int i = 0; i < 2; i++) {
+ acc += (make_fp32(va[i]) * make_fp32(vb[i]));
+ }
+ *((float *)(C_mem + m * N + n)) = acc;
+ }
+ }
+}
+
+int main() {
+ for (int i = 0; i < MATRIX_M; i++) {
+ for (int j = 0; j < MATRIX_K; j++) {
+ A[i][j] = bfloat16(1.0f * (i + j));
+ }
+ }
+ for (int i = 0; i < MATRIX_K / 2; i++) {
+ for (int j = 0; j < MATRIX_N * 2; j++) {
+ B[i][j] = bfloat16(2.0f * i + 3.0f * j);
+ }
+ }
+ for (int i = 0; i < MATRIX_M; i++) {
+ for (int j = 0; j < MATRIX_N; j++) {
+ C[i][j] = 1.0;
+ D[i][j] = 1.0;
+ }
+ }
+
+ std::cout << "M" << MATRIX_M << "N" << MATRIX_N << "K" << MATRIX_K << std::endl;
+
+ big_matrix MC((float *)&C);
+ big_matrix MD((float *)&D);
+ big_matrix MA((bfloat16 *)&A);
+ big_matrix MB((bfloat16 *)&B);
+
+ matrix_multiply(MC, MA, MB);
+
+ /*start = std::chrono::steady_clock::now();
+ matrix_multiply_ref((int32_t *)A, (int32_t *)B, (int32_t *)D, MATRIX_M, MATRIX_N, MATRIX_K / 2);
+ end = std::chrono::steady_clock::now();
+ std::cout << "Elapsed time in milliseconds (reference): " << std::chrono::duration_cast(end - start).count() << " ms" << std::endl;
+
+ bool res = true;
+ for (int i = 0; i < MATRIX_M; i++) {
+ for (int j = 0; j < MATRIX_N; j++) {
+ if ((fabs(C[i][j]) - fabs(D[i][j])) > BF16_EPSILON)
+ res = false;
+ }
+ }
+ std::cout << (res ? "passed" : "failed") << std::endl;
+ return !res;*/
+
+ return 0;
+}
+
diff --git a/extra/onnx.py b/extra/onnx.py
index ace358665a..61a557bb4b 100644
--- a/extra/onnx.py
+++ b/extra/onnx.py
@@ -181,7 +181,7 @@ def get_run_onnx(onnx_model: ModelProto):
fxn = getattr(onnx_ops, n.op_type)
if isinstance(fxn, dict):
for k in sorted(fxn.keys()):
- if k < onnx_model_version:
+ if k <= onnx_model_version:
real_fxn = fxn[k]
else:
real_fxn = fxn
diff --git a/extra/onnx_ops.py b/extra/onnx_ops.py
index 1ade5251d8..7db9199ef1 100644
--- a/extra/onnx_ops.py
+++ b/extra/onnx_ops.py
@@ -152,8 +152,10 @@ def Softmax_1(input, axis=1): return input.softmax(axis)
def Softmax_13(input, axis=-1): return input.softmax(axis)
Softmax = {1: Softmax_1, 13: Softmax_13} # Softmax default axis changed
def LogSoftmax(input, axis=-1): return input.log_softmax(axis)
-def Clip(input, min=-3.4e38, max=3.4e38): return input.clip(min, max)
-
+def Clip(input, min=None, max=None):
+ if min is None: min = -3.4e38
+ if max is None: max = 3.4e38
+ return input.clip(min, max)
def Sin(x): return x.sin()
def Cos(x): return x.cos()
@@ -169,7 +171,7 @@ def GreaterOrEqual(x:Tensor,y:Tensor): return (x>=y).cast(dtypes.bool)
def Equal(x:Tensor,y:Tensor): return (x==y).cast(dtypes.bool)
def Max(*data_0): return functools.reduce(Tensor.maximum, data_0)
-def Min(*data_0): return -functools.reduce(Tensor.maximum, [-x for x in data_0])
+def Min(*data_0): return functools.reduce(Tensor.minimum, data_0)
def Sum(*data_0): return functools.reduce(Tensor.__add__, data_0)
def Mean(*data_0): return functools.reduce(Tensor.__add__, data_0) / len(data_0)
diff --git a/test/external/external_test_onnx_backend.py b/test/external/external_test_onnx_backend.py
index 0799744558..8fc42cc619 100644
--- a/test/external/external_test_onnx_backend.py
+++ b/test/external/external_test_onnx_backend.py
@@ -60,7 +60,7 @@ backend_test.exclude('uint64')
backend_test.exclude('int8')
backend_test.exclude('int16')
backend_test.exclude('float64')
-
+backend_test.exclude('string')
backend_test.exclude('test_pow_types_int*')
backend_test.exclude('test_cast_*')
diff --git a/test/test_ops.py b/test/test_ops.py
index d0594d777e..28ba1bb40e 100644
--- a/test/test_ops.py
+++ b/test/test_ops.py
@@ -150,6 +150,10 @@ class TestOps(unittest.TestCase):
def test_mul(self):
helper_test_op([(64,64), (64,64)], lambda x,y: x*y, Tensor.mul)
helper_test_op([(), ()], lambda x,y: x*y, Tensor.mul)
+ def test_mul_const(self):
+ helper_test_op([(45,65)], lambda x: x*float("inf"), lambda x: x*float("inf"))
+ helper_test_op([(45,65)], lambda x: x*-float("inf"), lambda x: x*-float("inf"))
+ helper_test_op([(45,65)], lambda x: x*float("nan"), lambda x: x*float("nan"))
def test_div(self):
helper_test_op([(45,65), (45,65)], lambda x,y: x/y, Tensor.div)
helper_test_op([(), ()], lambda x,y: x/y, Tensor.div)
@@ -159,6 +163,12 @@ class TestOps(unittest.TestCase):
helper_test_op([(45,65)], lambda x: 1/x, lambda x: 1/x)
helper_test_op([(45,65)], lambda x: x/2, lambda x: x/2)
helper_test_op([(45,65)], lambda x: 2/x, lambda x: 2/x)
+ helper_test_op([(45,65)], lambda x: x/float("inf"), lambda x: x/float("inf"))
+ helper_test_op([(45,65)], lambda x: x/-float("inf"), lambda x: x/-float("inf"))
+ helper_test_op([(45,65)], lambda x: x/float("nan"), lambda x: x/float("nan"))
+ helper_test_op([(45,65)], lambda x: float("inf")/x, lambda x: float("inf")/x)
+ helper_test_op([(45,65)], lambda x: (-float("inf"))/x, lambda x: (-float("inf"))/x)
+ helper_test_op([(45,65)], lambda x: float("nan")/x, lambda x: float("nan")/x)
helper_test_op([()], lambda x: x/2, lambda x: x/2)
helper_test_op([()], lambda x: 2/x, lambda x: 2/x)
def test_pow(self):
@@ -247,6 +257,12 @@ class TestOps(unittest.TestCase):
with self.assertRaises(RuntimeError):
a = Tensor(3.14)
a.matmul(a)
+
+ def test_cumsum(self):
+ helper_test_op([(20)], lambda x: torch.cumsum(x, dim=0), lambda x: Tensor.cumsum(x, axis=0), atol=1e-6)
+ helper_test_op([(20,30)], lambda x: torch.cumsum(x, dim=0), lambda x: Tensor.cumsum(x, axis=0), atol=1e-6)
+ helper_test_op([(20,30)], lambda x: torch.cumsum(x, dim=1), lambda x: Tensor.cumsum(x, axis=1), atol=1e-6)
+ helper_test_op([(20,30,40)], lambda x: torch.cumsum(x, dim=2), lambda x: Tensor.cumsum(x, axis=2), atol=1e-6)
def test_matmul_simple(self):
helper_test_op([(4), (4,4)], lambda x,y: x.matmul(y), Tensor.dot, atol=1e-4)
def test_matmul(self):
@@ -451,6 +467,13 @@ class TestOps(unittest.TestCase):
a[1, 77] # IndexError: (out of bounds).
a[0, -77]
+ def test_slice_ellipsis(self):
+ helper_test_op([(3,3,3,3)], lambda x: x[..., 0], lambda x: x[..., 0])
+ helper_test_op([(3,3,3,3)], lambda x: x[0, ...], lambda x: x[0, ...])
+ helper_test_op([(3,3,3,3)], lambda x: x[0, ..., 0], lambda x: x[0, ..., 0])
+ helper_test_op([(3,3,3,3)], lambda x: x[0:3, ..., 2:3], lambda x: x[0:3, ..., 2:3])
+ helper_test_op([(3,3,3,3)], lambda x: x[None, 0:3, ..., 0, None], lambda x: x[None, 0:3, ..., 0, None])
+
def test_pad2d(self):
helper_test_op([(3,3,3,3)], lambda x: torch.nn.functional.pad(x, (1,2,3,4)), lambda x: x.pad2d(padding=(1,2,3,4)))
diff --git a/tinygrad/codegen/cstyle.py b/tinygrad/codegen/cstyle.py
index 73de73b159..0633986760 100644
--- a/tinygrad/codegen/cstyle.py
+++ b/tinygrad/codegen/cstyle.py
@@ -120,8 +120,10 @@ def uops_to_cstyle(uops:List[UOp], bufs:List[Union[LocalBuffer,LazyBuffer]], lan
# TODO: merge with CONST?
if bufs[args.i] is not None and isinstance(bufs[args.i].realized, RawConst):
assert newvar.ltype == LocalTypes.float, "const can't be float4"
- # nan? inf?
- val = f"{bufs[args.i].realized._buf}" + ("f" if not dtypes.is_int(bufs[args.i].dtype) else "")
+ x = bufs[args.i].realized._buf
+ if math.isnan(x): val = "NAN"
+ elif math.isinf(x): val = ("-" if x < 0 else "") + "INFINITY"
+ else: val = f"{x}" + ("f" if not dtypes.is_int(bufs[args.i].dtype) else "")
elif isinstance(bufs[args.i].dtype, ImageDType):
assert newvar.ltype == LocalTypes.float4, "image must be float4"
prekernel.add("const sampler_t smp = CLK_NORMALIZED_COORDS_FALSE | CLK_ADDRESS_CLAMP | CLK_FILTER_NEAREST;\n")
diff --git a/tinygrad/lazy.py b/tinygrad/lazy.py
index 74bbc2bade..b17988c560 100644
--- a/tinygrad/lazy.py
+++ b/tinygrad/lazy.py
@@ -7,6 +7,7 @@ from tinygrad.helpers import prod, getenv, DType, dtypes, flatten, ImageDType, D
from tinygrad.shape.shapetracker import ShapeTracker, get_contraction
from tinygrad.ops import Compiled, Interpreted, UnaryOps, BinaryOps, ReduceOps, MovementOps, LoadOps, OpType, LazyOp, get_lazyops, get_buffers, map_buffers
from tinygrad.runtime.lib import RawConst, RawBuffer, RawBufferMapped
+from tinygrad.runtime.ops_cpu import RawNumpyBuffer
from tinygrad.runtime.ops_disk import RawDiskBuffer
# lazy can recurse a lot
@@ -73,7 +74,7 @@ def create_lazybuffer(device:str, shape:Union[ShapeTracker, Tuple[int, ...]], op
st = shape if isinstance(shape, ShapeTracker) else ShapeTracker(tuple(shape))
# fromcpu aren't cached
- if optype == LoadOps and op.op in [LoadOps.FROMCPU, LoadOps.EMPTY, LoadOps.RAND, LoadOps.CONST]: return LazyBuffer(device, st, optype, op, dtype)
+ if optype == LoadOps and op.op in [LoadOps.EMPTY, LoadOps.RAND, LoadOps.CONST]: return LazyBuffer(device, st, optype, op, dtype)
#print("create_lazybuffer", device, shape, optype, op, dtype)
@@ -87,16 +88,17 @@ def create_lazybuffer(device:str, shape:Union[ShapeTracker, Tuple[int, ...]], op
class LazyBuffer:
__deletable__ = ('op',)
- def __init__(self, device:str, st:ShapeTracker, optype:OpType, op:LazyOp, dtype:DType):
+ def __init__(self, device:str, st:ShapeTracker, optype:OpType, src:Union[LazyOp, RawBuffer], dtype:DType):
self.st = st # NOTE: this is not a copy! this should be a "read-only" ShapeTracker
self.device, self.shape, self.optype, self.dtype = device, self.st.shape, optype, dtype
- self.op: LazyOp = op
- self.realized: Optional[RawBuffer] = None
+ self.realized: Optional[RawBuffer] = src if isinstance(src, RawBuffer) else None
self.output_buffer: Optional[RawBuffer] = None # TODO: do we really need this? or can we just use realized
# TODO: does children have to be a ref count instead of a set? can a Buffer be a double child?
self.children: weakref.WeakSet[LazyBuffer] = weakref.WeakSet()
# NOTE: op should be read only after construction of LazyBuffer
- for x in get_buffers(op): x.children.add(self)
+ if isinstance(src, LazyOp):
+ self.op: LazyOp = src
+ for x in get_buffers(self.op): x.children.add(self)
if not LAZY: self.realize()
# log phantom ops to the graph
@@ -109,10 +111,7 @@ class LazyBuffer:
def realize(self:LazyBuffer) -> LazyBuffer:
if self.realized is None:
# get real ops first
- if self.op.op == LoadOps.FROMCPU:
- if DEBUG >= 4: print(f"copying {self.op.arg.shape}:{dtypes.from_np(self.op.arg.dtype)} -> {self.device}")
- self.realized = Device[self.device].buffer.fromCPU(self.op.arg, **self._device_extra_args())
- elif self.op.op == LoadOps.CONTIGUOUS:
+ if self.op.op == LoadOps.CONTIGUOUS:
realized = self.op.src[0].realize().realized
if self.op.src[0].st.contiguous and not isinstance(realized, RawConst) and realized.size == prod(self.shape):
# no need to run an AST, this is already contiguous
@@ -179,6 +178,10 @@ class LazyBuffer:
def loadop(op, shape, dtype, device, arg=None, src=None) -> LazyBuffer:
return create_lazybuffer(device, shape, LoadOps, LazyOp(op, tuple() if src is None else (src,), arg), dtype)
+ @staticmethod
+ def fromCPU(x: np.ndarray) -> LazyBuffer:
+ return LazyBuffer("CPU", ShapeTracker(x.shape), LoadOps, RawNumpyBuffer.fromCPU(x), dtypes.from_np(x.dtype))
+
# create a constant with the shape and dtype of self
def const_like(self, val) -> LazyBuffer:
# NOTE: dtypes.from_np(self.dtype.np) to deal with image types
diff --git a/tinygrad/nn/__init__.py b/tinygrad/nn/__init__.py
index 5f1c7336a0..65dbd2b9c1 100644
--- a/tinygrad/nn/__init__.py
+++ b/tinygrad/nn/__init__.py
@@ -35,6 +35,10 @@ class BatchNorm2d:
return x.batchnorm(self.weight, self.bias, batch_mean, batch_invstd)
+# TODO: these Conv lines are terrible
+def Conv1d(in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True):
+ return Conv2d(in_channels, out_channels, (kernel_size,), stride, padding, dilation, groups, bias)
+
class Conv2d:
def __init__(self, in_channels, out_channels, kernel_size, stride=1, padding=0, dilation=1, groups=1, bias=True):
self.kernel_size = (kernel_size, kernel_size) if isinstance(kernel_size, int) else tuple(kernel_size)
diff --git a/tinygrad/ops.py b/tinygrad/ops.py
index f2cd5e0a49..720b22cafd 100644
--- a/tinygrad/ops.py
+++ b/tinygrad/ops.py
@@ -12,7 +12,7 @@ class UnaryOps(Enum): NOOP = auto(); EXP = auto(); LOG = auto(); CAST = auto();
class BinaryOps(Enum): ADD = auto(); SUB = auto(); MUL = auto(); DIV = auto(); POW = auto(); CMPEQ = auto(); MAX = auto() # noqa: E702
class ReduceOps(Enum): SUM = auto(); MAX = auto() # noqa: E702
class FusedOps(Enum): MULACC = auto() # noqa: E702
-class LoadOps(Enum): EMPTY = auto(); RAND = auto(); CONST = auto(); FROM = auto(); FROMCPU = auto(); CONTIGUOUS = auto(); CUSTOM = auto() # noqa: E702
+class LoadOps(Enum): EMPTY = auto(); RAND = auto(); CONST = auto(); FROM = auto(); CONTIGUOUS = auto(); CUSTOM = auto() # noqa: E702
Op = Union[UnaryOps, BinaryOps, ReduceOps, MovementOps, LoadOps, FusedOps]
OpType = Union[Type[UnaryOps], Type[BinaryOps], Type[ReduceOps], Type[MovementOps], Type[LoadOps], Type[FusedOps]]
diff --git a/tinygrad/runtime/lib.py b/tinygrad/runtime/lib.py
index d29d81b006..313e13b0a4 100644
--- a/tinygrad/runtime/lib.py
+++ b/tinygrad/runtime/lib.py
@@ -36,7 +36,7 @@ class RawBufferMapped(RawBufferCopyIn):
# this one is simple enough that i moved it out of the runtimes
class RawMallocBuffer(RawBufferMapped):
- def __init__(self, size, dtype: DType): super().__init__(size, dtype, ({dtypes.float32: ctypes.c_float, dtypes.float16: ctypes.c_int16, dtypes.int8: ctypes.c_int8, dtypes.uint8: ctypes.c_uint8, dtypes.bool: ctypes.c_uint8, dtypes.int64: ctypes.c_int64}[dtype] * size)())
+ def __init__(self, size, dtype: DType): super().__init__(size, dtype, ({dtypes.float32: ctypes.c_float, dtypes.float16: ctypes.c_int16, dtypes.int8: ctypes.c_int8, dtypes.uint8: ctypes.c_uint8, dtypes.bool: ctypes.c_uint8, dtypes.int32: ctypes.c_int32, dtypes.int64: ctypes.c_int64}[dtype] * size)())
def _buffer(self): return memoryview(self._buf)
class RawBufferCopyInOut(RawBufferCopyIn):
diff --git a/tinygrad/state.py b/tinygrad/state.py
index 567c058b92..13d7d35d71 100644
--- a/tinygrad/state.py
+++ b/tinygrad/state.py
@@ -45,10 +45,12 @@ def get_parameters(obj) -> List[Tensor]: return list(get_state_dict(obj).values(
def load_state_dict(model, state_dict, strict=True):
with Timing("loaded weights in ", lambda et_ns: f", {GlobalCounters.mem_used/1e9:.2f} GB loaded at {GlobalCounters.mem_used/et_ns:.2f} GB/s"):
- for k,v in (t := tqdm(get_state_dict(model).items())):
+ model_state_dict = get_state_dict(model)
+ if DEBUG >= 1 and len(state_dict) > len(model_state_dict): print("WARNING: unused weights in state_dict", sorted(list(state_dict.keys() - model_state_dict.keys())))
+ for k,v in (t := tqdm(model_state_dict.items())):
t.set_description(f"ram used: {GlobalCounters.mem_used/1e9:5.2f} GB, {k:50s}")
if k not in state_dict and not strict:
- if DEBUG >= 2: print(f"WARNING: not loading {k}")
+ if DEBUG >= 1: print(f"WARNING: not loading {k}")
continue
v.assign(state_dict[k].to(v.device)).realize()
@@ -72,7 +74,7 @@ def torch_load(fn:str):
if tuple(permute_indexes) != tuple(range(len(permute_indexes))):
intermediate_shape = tuple([shape_strides[x][0] for x in argsort(permute_indexes)])
assert tuple([shape_strides[i][1] for i in argsort(permute_indexes)]) == strides_for_shape(intermediate_shape), "nonpermutable strides"
- if DEBUG >= 2: print(f"WARNING: this torch load is slow. it has to convert to CPU to permute {permute_indexes}")
+ if DEBUG >= 2: print(f"WARNING: this torch load is slow. CPU to permute {intermediate_shape} with {permute_indexes}")
# TODO: find a nice way to support all shapetracker on disktensors
ret = ret.cpu().reshape(intermediate_shape).permute(permute_indexes)
diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py
index 239710552d..78d5828353 100644
--- a/tinygrad/tensor.py
+++ b/tinygrad/tensor.py
@@ -3,7 +3,7 @@ from __future__ import annotations
import math, functools, itertools, operator
import numpy as np
from typing import List, Tuple, Callable, Optional, ClassVar, Type, Union, Sequence
-from tinygrad.helpers import prod, argfix, make_pair, getenv, IMAGE, DEBUG, flatten, DType, dtypes
+from tinygrad.helpers import prod, argfix, make_pair, getenv, IMAGE, DEBUG, flatten, DType, dtypes, ImageDType
from tinygrad.lazy import Device, LazyBuffer
from tinygrad.ops import LoadOps
@@ -34,18 +34,17 @@ class Tensor:
no_grad: ClassVar[bool] = False
default_type: ClassVar[DType] = dtypes.float32
- def __init__(self, data:Union[int, float, list, LazyBuffer, np.ndarray], device=Device.DEFAULT, dtype:Optional[DType]=None, requires_grad:Optional[bool]=None):
+ def __init__(self, data:Union[int, float, list, tuple, LazyBuffer, np.ndarray], device=Device.DEFAULT, dtype:Optional[DType]=None, requires_grad:Optional[bool]=None):
assert dtype is None or isinstance(dtype, DType), f"invalid dtype {dtype}"
device = Device.canonicalize(device)
- if isinstance(data, list):
+ if isinstance(data, (list, tuple)):
data = np.array(data, dtype=(dtype if dtype is not None else Tensor.default_type).np)
+ if isinstance(data, np.ndarray):
+ data = LazyBuffer.fromCPU(data)
if isinstance(data, LazyBuffer):
assert dtype is None or dtype == data.dtype, "dtype doesn't match, and casting isn't supported"
lazydata = data if data.device == device else LazyBuffer.loadop(LoadOps.FROM, data.shape, data.dtype, device, src=data)
- elif isinstance(data, np.ndarray):
- # TODO: create CPUBuffer directly
- lazydata = LazyBuffer.loadop(LoadOps.FROMCPU, data.shape, dtypes.from_np(data.dtype), device, data)
elif isinstance(data, (int, float)):
lazydata = LazyBuffer.loadop(LoadOps.CONST, tuple(), dtype if dtype is not None else Tensor.default_type, device, data)
else:
@@ -65,7 +64,7 @@ class Tensor:
self._ctx: Optional[Function] = None
def __repr__(self):
- return f""
+ return f""
# Python has a non moving GC, so this should be okay
def __hash__(self): return id(self)
@@ -263,7 +262,15 @@ class Tensor:
val = list(val) if isinstance(val, tuple) else [val]
if (num_slices := sum(isinstance(v, (slice, int)) for v in val)) > len(self.shape):
raise IndexError(f"too many indices for tensor of dimension {len(self.shape)}")
- orig_slices = list(val) + [slice(None)] * (len(self.shape) - num_slices)
+ orig_slices = list(val)
+ ellipses_found = [i for i, v in enumerate(val) if v is Ellipsis]
+ if len(ellipses_found) > 0:
+ if len(ellipses_found) != 1:
+ raise IndexError("an index can only have a single ellipsis ('...')")
+ ellipsis_idx = ellipses_found[0]
+ orig_slices[ellipsis_idx:ellipsis_idx+1] = [slice(None)] * (len(self.shape) - num_slices)
+ else:
+ orig_slices += [slice(None)] * (len(self.shape) - num_slices)
valid_slices = list(itertools.filterfalse(lambda x: x is None, orig_slices))
valid_slices = [v if isinstance(v, slice) else slice(y := normalize_int(v, i, dim_sz), y+1) for i, (v, dim_sz) in enumerate(zip(valid_slices, self.shape))]
start, stop, strides = zip(*y) if (y := [s.indices(dim_sz) for s, dim_sz in zip(valid_slices, self.shape)]) else ((), (), ())
@@ -484,9 +491,10 @@ class Tensor:
r = (x*w).sum(-1)
return r.reshape((*r.shape[:-2], r.shape[-1])) if len(self.shape) == 1 else r
- # TODO: make this work for n-dimensional inputs
- def cumsum(self): return self.reshape(1, 1, 1, self.shape[0]).conv2d(Tensor.ones(1, 1, 1, self.shape[0]), padding=(self.shape[0] - 1, 0, 0, 0)).flatten()
-
+ def cumsum(self, axis=0):
+ x = self.permute(*(i for i in range(self.ndim) if i != axis), axis)
+ return x.reshape(1, 1, -1, self.shape[axis]).conv2d(Tensor.ones(1, 1, 1, self.shape[axis]), padding=(self.shape[axis]-1, 0, 0, 0)).reshape(*x.shape).permute(*range(axis), self.ndim - 1, *range(axis, self.ndim-1))
+
# ***** mlops (unary) *****
def contiguous(self): return mlops.Contiguous.apply(self)
@@ -503,7 +511,7 @@ class Tensor:
def sqrt(self): return self.pow(0.5)
def rsqrt(self): return self.pow(-0.5)
def square(self): return self*self
- def clip(self, min_, max_): return ((self-min_).relu()+min_) - (self-max_).relu()
+ def clip(self, min_, max_): return self.maximum(min_).minimum(max_)
def abs(self): return self.relu() + (-self).relu()
def sign(self): return self / (self.abs() + 1e-10)
def reciprocal(self): return 1.0/self
@@ -527,9 +535,9 @@ class Tensor:
def softsign(self): return self / (1 + self.abs())
# ***** broadcasted binary mlops *****
-
def _broadcasted(self, fxn:Type[Function], other:Union[Tensor, float], reverse:bool=False) -> Tensor:
- x,y = [Tensor(t, device=self.device, requires_grad=False) if not isinstance(t, Tensor) else t for t in ([other,self] if reverse else [self,other])]
+ dtype = self.dtype if self.dtype != dtypes.bool and not isinstance(self.dtype,ImageDType) else dtypes.float32
+ x,y = [Tensor(t, device=self.device, requires_grad=False, dtype=dtype) if not isinstance(t, Tensor) else t for t in ([other,self] if reverse else [self,other])]
x,y = [t.reshape([1]*(max(len(x.shape), len(y.shape))-len(t.shape)) + list(t.shape)) for t in [x,y]]
shape_ret = tuple(max(sx, sy) for sx,sy in zip(x.shape, y.shape))
return fxn.apply(x.expand(shape_ret), y.expand(shape_ret))