Tinygrad in browser

Require JavaScript Promise Integration (JSPI) to work. You can find out more about it here and track its adoption here.

To view or modify the setting, enter this in the browser URL: chrome://flags/#enable-experimental-webassembly-jspi

JSPI status:

On Ubuntu, you will need to launch Chrome in the CLI with: google-chrome --enable-features=Vulkan --enable-unsafe-webgpu

from typing import List, Callable from tinygrad import Tensor, TinyJit, nn, GlobalCounters from tinygrad.helpers import getenv, colored, trange, DEBUG DEBUG.value = 2 a = Tensor.arange(16).reshape((4, 4)) print(a.tolist()) class Model: def __init__(self): self.layers: List[Callable[[Tensor], Tensor]] = [ nn.Conv2d(1, 32, 5), Tensor.relu, nn.Conv2d(32, 32, 5), Tensor.relu, nn.BatchNorm(32), Tensor.max_pool2d, nn.Conv2d(32, 64, 3), Tensor.relu, nn.Conv2d(64, 64, 3), Tensor.relu, nn.BatchNorm(64), Tensor.max_pool2d, lambda x: x.flatten(1), nn.Linear(576, 10)] def __call__(self, x:Tensor) -> Tensor: return x.sequential(self.layers) model = Model() opt = nn.optim.Adam(nn.state.get_parameters(model)) @Tensor.train() def train(): batch = 1024 x = Tensor.rand(batch, 1, 28, 28) y = Tensor.rand(batch) loss = model(x).sparse_categorical_crossentropy(y).backward() opt.step() print(f"{loss.tolist()=}") train()