IQ.Pilot Release Commit @ bec7652

artifacts/package_sources/tinygrad/test/models/test_bert.py

@@ -0,0 +1,58 @@
+#!/usr/bin/env python
+import os
+os.environ['USE_TF'] = '0'  # prevent transformers from importing tensorflow
+import unittest
+from tinygrad import Tensor
+import numpy as np
+import torch
+
+def get_question_samp(bsz, seq_len, vocab_size, seed):
+  np.random.seed(seed)
+  in_ids = np.random.randint(vocab_size, size=(bsz, seq_len), dtype=np.int32)
+  mask = np.random.choice([True, False], size=(bsz, seq_len))
+  seg_ids = np.random.randint(2, size=(bsz, seq_len), dtype=np.int32)  # type_vocab_size
+  return in_ids, mask, seg_ids
+
+def set_equal_weights(mdl, torch_mdl):
+  from tinygrad.nn.state import get_state_dict
+  state, torch_state = get_state_dict(mdl), torch_mdl.state_dict()
+  assert len(state) == len(torch_state)
+  for k, v in state.items():
+    assert k in torch_state
+    torch_state[k].copy_(torch.from_numpy(v.numpy()))
+  torch_mdl.eval()
+
+class TestBert(unittest.TestCase):
+  def test_questions(self):
+    from extra.models.bert import BertForQuestionAnswering
+    from transformers import BertForQuestionAnswering as TorchBertForQuestionAnswering
+    from transformers import BertConfig
+
+    # small
+    config = {
+      'vocab_size':24, 'hidden_size':2, 'num_hidden_layers':2, 'num_attention_heads':2,
+      'intermediate_size':32, 'hidden_dropout_prob':0.1, 'attention_probs_dropout_prob':0.1,
+      'max_position_embeddings':512, 'type_vocab_size':2
+      }
+
+    # Create in tinygrad
+    Tensor.manual_seed(1337)
+    mdl = BertForQuestionAnswering(**config)
+
+    # Create in torch
+    with torch.no_grad():
+      torch_mdl = TorchBertForQuestionAnswering(BertConfig(**config))
+
+    set_equal_weights(mdl, torch_mdl)
+
+    seeds = (1337, 3141)
+    bsz, seq_len = 1, 16
+    for seed in seeds:
+      in_ids, mask, seg_ids = get_question_samp(bsz, seq_len, config['vocab_size'], seed)
+      out = mdl(Tensor(in_ids), Tensor(mask), Tensor(seg_ids))
+      torch_out = torch_mdl.forward(torch.from_numpy(in_ids).long(), torch.from_numpy(mask), torch.from_numpy(seg_ids).long())[:2]
+      torch_out = torch.cat(torch_out).unsqueeze(2)
+      np.testing.assert_allclose(out.numpy(), torch_out.detach().numpy(), atol=5e-4, rtol=5e-4)
+
+if __name__ == '__main__':
+  unittest.main()

artifacts/package_sources/tinygrad/test/models/test_efficientnet.py

@@ -0,0 +1,114 @@
+import ast, pathlib, unittest
+
+import numpy as np
+from PIL import Image
+
+from tinygrad import Tensor, Context
+from tinygrad.helpers import getenv
+from test.helpers import slow
+from extra.models.efficientnet import EfficientNet
+from extra.models.vit import ViT
+from extra.models.resnet import ResNet50
+
+def _load_labels():
+  labels_filename = pathlib.Path(__file__).parent / 'efficientnet/imagenet1000_clsidx_to_labels.txt'
+  return ast.literal_eval(labels_filename.read_text())
+
+_LABELS = _load_labels()
+
+def preprocess(img, new=False):
+  # preprocess image
+  aspect_ratio = img.size[0] / img.size[1]
+  img = img.resize((int(224*max(aspect_ratio,1.0)), int(224*max(1.0/aspect_ratio,1.0))))
+
+  img = np.array(img)
+  y0, x0 =(np.asarray(img.shape)[:2] - 224) // 2
+  img = img[y0: y0 + 224, x0: x0 + 224]
+
+  # low level preprocess
+  if new:
+    img = img.astype(np.float32)
+    img -= [127.0, 127.0, 127.0]
+    img /= [128.0, 128.0, 128.0]
+    img = img[None]
+  else:
+    img = np.moveaxis(img, [2, 0, 1], [0, 1, 2])
+    img = img.astype(np.float32)[:3].reshape(1, 3, 224, 224)
+    img /= 255.0
+    img -= np.array([0.485, 0.456, 0.406]).reshape((1, -1, 1, 1))
+    img /= np.array([0.229, 0.224, 0.225]).reshape((1, -1, 1, 1))
+  return img
+
+def _infer(model: EfficientNet, img):
+  with Context(TRAINING=0):
+    out = model.forward(Tensor(img)).argmax(axis=-1)
+  return out.tolist()
+
+chicken_img = preprocess(Image.open(pathlib.Path(__file__).parent / 'efficientnet/Chicken.jpg'))
+car_img = preprocess(Image.open(pathlib.Path(__file__).parent / 'efficientnet/car.jpg'))
+
+class TestEfficientNet(unittest.TestCase):
+  @classmethod
+  def setUpClass(cls):
+    cls.model = EfficientNet(number=getenv("NUM"))
+    cls.model.load_from_pretrained()
+
+  @classmethod
+  def tearDownClass(cls):
+    del cls.model
+
+  @slow
+  def test_chicken(self):
+    labels = _infer(self.model, chicken_img)
+    self.assertEqual(_LABELS[labels[0]], "hen")
+
+  @slow
+  def test_car(self):
+    labels = _infer(self.model, car_img)
+    self.assertEqual(_LABELS[labels[0]], "sports car, sport car")
+
+  def test_chicken_car(self):
+    labels = _infer(self.model, np.concatenate([chicken_img, car_img], axis=0))
+    self.assertEqual(_LABELS[labels[0]], "hen")
+    self.assertEqual(_LABELS[labels[1]], "sports car, sport car")
+
+@unittest.skip("these pretrained models are no longer available")
+class TestViT(unittest.TestCase):
+  @classmethod
+  def setUpClass(cls):
+    cls.model = ViT()
+    cls.model.load_from_pretrained()
+
+  @classmethod
+  def tearDownClass(cls):
+    del cls.model
+
+  def test_chicken(self):
+    labels = _infer(self.model, chicken_img)
+    self.assertEqual(_LABELS[labels[0]], "cock")
+
+  def test_car(self):
+    labels = _infer(self.model, car_img)
+    self.assertEqual(_LABELS[labels[0]], "racer, race car, racing car")
+
+class TestResNet(unittest.TestCase):
+  @classmethod
+  def setUpClass(cls):
+    cls.model = ResNet50()
+    cls.model.load_from_pretrained()
+
+  @classmethod
+  def tearDownClass(cls):
+    del cls.model
+
+  def test_chicken(self):
+    labels = _infer(self.model, chicken_img)
+    # NOTE: logits for these two are close
+    self.assertIn(_LABELS[labels[0]], ("hen", "cock"))
+
+  def test_car(self):
+    labels = _infer(self.model, car_img)
+    self.assertEqual(_LABELS[labels[0]], "sports car, sport car")
+
+if __name__ == '__main__':
+  unittest.main()

artifacts/package_sources/tinygrad/test/models/test_end2end.py

@@ -0,0 +1,164 @@
+import torch
+from torch import nn
+import unittest, sys
+import numpy as np
+from tinygrad.nn.state import get_parameters, get_state_dict
+from tinygrad.nn import optim, Linear, Conv2d, BatchNorm2d
+from tinygrad.tensor import Tensor
+from tinygrad.helpers import Context
+from extra.datasets import fetch_mnist
+
+def compare_tiny_torch(model, model_torch, X, Y):
+  with Context(TRAINING=1):
+    model_torch.train()
+    model_state_dict = get_state_dict(model)
+    for k,v in model_torch.named_parameters():
+      if sys.stdout.isatty(): print(f"initting {k} from torch")
+      model_state_dict[k].assign(Tensor(v.detach().numpy())).realize()
+
+    optimizer = optim.SGD(get_parameters(model), lr=0.001)
+    optimizer_torch = torch.optim.SGD(model_torch.parameters(), lr=0.001)
+
+    Xt = torch.Tensor(X.numpy())
+    np.testing.assert_allclose(X.numpy(), Xt.detach().numpy())
+
+    out = model(X)
+    loss = (out * Y).mean()
+
+    out_torch = model_torch(torch.Tensor(X.numpy()))
+    loss_torch = (out_torch * torch.Tensor(Y.numpy())).mean()
+
+    # zero and backward
+    optimizer.zero_grad()
+    loss.backward()
+    optimizer_torch.zero_grad()
+    loss_torch.backward()
+
+    # assert losses match
+    if sys.stdout.isatty(): print(loss.realize().numpy())
+    if sys.stdout.isatty(): print(loss_torch.detach().numpy())
+    np.testing.assert_allclose(loss.realize().numpy(), loss_torch.detach().numpy(), atol=1e-4)
+
+    for k,v in list(model_torch.named_parameters())[::-1]:
+      g = model_state_dict[k].grad.numpy()
+      gt = v.grad.detach().numpy()
+      if sys.stdout.isatty(): print("testing grads", k, model_state_dict[k].grad.dtype)
+      np.testing.assert_allclose(g, gt, atol=1e-3, err_msg=f'grad mismatch {k}')
+
+    # take the steps
+    optimizer.step()
+    optimizer_torch.step()
+
+    # assert weights match
+    for k,v in model_torch.named_parameters():
+      if sys.stdout.isatty(): print("testing weight", k, model_state_dict[k].dtype)
+      np.testing.assert_allclose(model_state_dict[k].numpy(), v.detach().numpy(), atol=1e-3, err_msg=f'weight mismatch {k}')
+
+def get_mnist_data():
+  _X_train, _Y_train, X_test, Y_test = fetch_mnist()
+  BS = 32
+  num_classes = 10
+  X = Tensor(X_test[0:BS].astype(np.float32))
+  Y = np.zeros((BS, num_classes), np.float32)
+  Y[range(BS),Y_test[0:BS]] = -1.0*num_classes
+  return X, Tensor(Y)
+
+class TestEnd2End(unittest.TestCase):
+  @classmethod
+  def setUpClass(cls):
+    cls.X, cls.Y = get_mnist_data()
+
+  def setUp(self):
+    torch.manual_seed(123)
+
+  def test_linear_mnist(self):
+    class LinTiny:
+      def __init__(self, bias=False):
+        self.l1 = Linear(784, 128, bias=bias)
+        self.l2 = Linear(128, 10, bias=bias)
+      def __call__(self, x):
+        return self.l2(self.l1(x).relu()).log_softmax(-1)
+    class LinTorch(nn.Module):
+      def __init__(self, bias=False):
+        super().__init__()
+        self.l1 = nn.Linear(784, 128, bias=bias)
+        self.l2 = nn.Linear(128, 10, bias=bias)
+      def forward(self, x):
+        return self.l2(self.l1(x).relu()).log_softmax(-1)
+    compare_tiny_torch(LinTiny(), LinTorch(), self.X, self.Y)
+
+  def test_bn_mnist(self):
+    class LinTiny:
+      def __init__(self):
+        self.l1 = Linear(784, 128)
+        self.l2 = Linear(128, 10)
+        self.bn1 = BatchNorm2d(128)
+      def __call__(self, x):
+        return self.l2(self.bn1(self.l1(x).reshape(x.shape[0], -1, 1, 1)).reshape(x.shape[0], -1).relu()).log_softmax(-1)
+    class LinTorch(nn.Module):
+      def __init__(self):
+        super().__init__()
+        self.l1 = nn.Linear(784, 128)
+        self.l2 = nn.Linear(128, 10)
+        self.bn1 = nn.BatchNorm2d(128)
+      def forward(self, x):
+        return self.l2(self.bn1(self.l1(x).reshape(x.shape[0], -1, 1, 1)).reshape(x.shape[0], -1).relu()).log_softmax(-1)
+    compare_tiny_torch(LinTiny(), LinTorch(), self.X, self.Y)
+
+  def test_bn_alone(self):
+    np.random.seed(1337)
+    X = Tensor(np.random.randn(32, 10, 1, 1).astype(np.float32))
+    Y = Tensor(np.random.randn(32, 10, 1, 1).astype(np.float32))
+    compare_tiny_torch(BatchNorm2d(10), nn.BatchNorm2d(10), X, Y)
+
+  def test_bn_linear(self):
+    BS, K = 2, 1
+    eps = 1e-12  # torch asserts if this is 0
+    X = Tensor([1,0]).reshape(BS, K, 1, 1)
+    Y = Tensor([-1,0]).reshape(BS, K, 1, 1)
+    class LinTiny:
+      def __init__(self):
+        self.l1 = Conv2d(K, K, 1, bias=False)
+        self.bn1 = BatchNorm2d(K, affine=False, track_running_stats=False, eps=eps)
+      def __call__(self, x): return self.bn1(self.l1(x))
+    class LinTorch(nn.Module):
+      def __init__(self):
+        super().__init__()
+        self.l1 = nn.Conv2d(K, K, 1, bias=False)
+        self.bn1 = nn.BatchNorm2d(K, affine=False, track_running_stats=False, eps=eps)
+      def forward(self, x): return self.bn1(self.l1(x))
+    model_torch = LinTorch()
+    with torch.no_grad():
+      model_torch.l1.weight[:] = 1.
+    compare_tiny_torch(LinTiny(), model_torch, X, Y)
+
+  def test_conv_mnist(self):
+    class LinTiny:
+      def __init__(self, has_batchnorm=False):
+        self.c1 = Conv2d(1, 8, 3, stride=2)
+        self.c2 = Conv2d(8, 16, 3, stride=2)
+        self.l1 = Linear(16*6*6, 10)
+        if has_batchnorm:
+          self.bn1, self.bn2 = BatchNorm2d(8), BatchNorm2d(16)
+        else:
+          self.bn1, self.bn2 = lambda x: x, lambda x: x
+      def __call__(self, x):
+        return self.l1(self.bn2(self.c2(self.bn1(self.c1(x)).relu())).relu().reshape(x.shape[0], -1)).log_softmax(-1)
+    class LinTorch(nn.Module):
+      def __init__(self, has_batchnorm=False):
+        super().__init__()
+        self.c1 = nn.Conv2d(1, 8, 3, stride=2)
+        self.c2 = nn.Conv2d(8, 16, 3, stride=2)
+        self.l1 = nn.Linear(16*6*6, 10)
+        if has_batchnorm:
+          self.bn1, self.bn2 = nn.BatchNorm2d(8), nn.BatchNorm2d(16)
+        else:
+          self.bn1, self.bn2 = lambda x: x, lambda x: x
+      def forward(self, x):
+        return self.l1(self.bn2(self.c2(self.bn1(self.c1(x)).relu())).relu().reshape(x.shape[0], -1)).log_softmax(-1)
+    for has_batchnorm in [False, True]:
+      with self.subTest(has_batchnorm=has_batchnorm):
+        compare_tiny_torch(LinTiny(has_batchnorm), LinTorch(has_batchnorm), self.X.reshape((-1, 1, 28, 28)), self.Y)
+
+if __name__ == "__main__":
+  unittest.main()

artifacts/package_sources/tinygrad/test/models/test_mnist.py

@@ -0,0 +1,115 @@
+#!/usr/bin/env python
+import unittest
+import numpy as np
+from tinygrad import Tensor
+from test.helpers import slow
+from tinygrad.nn.state import get_parameters
+from tinygrad.nn import optim, BatchNorm2d
+from extra.training import train, evaluate
+from extra.datasets import fetch_mnist
+
+# load the mnist dataset
+X_train, Y_train, X_test, Y_test = fetch_mnist()
+
+# create a model
+class TinyBobNet:
+  def __init__(self):
+    self.l1 = Tensor.scaled_uniform(784, 128)
+    self.l2 = Tensor.scaled_uniform(128, 10)
+
+  def parameters(self):
+    return get_parameters(self)
+
+  def forward(self, x):
+    return x.dot(self.l1).relu().dot(self.l2)
+
+# create a model with a conv layer
+class TinyConvNet:
+  def __init__(self, has_batchnorm=False):
+    # https://keras.io/examples/vision/mnist_convnet/
+    conv = 3
+    #inter_chan, out_chan = 32, 64
+    inter_chan, out_chan = 8, 16   # for speed
+    self.c1 = Tensor.scaled_uniform(inter_chan,1,conv,conv)
+    self.c2 = Tensor.scaled_uniform(out_chan,inter_chan,conv,conv)
+    self.l1 = Tensor.scaled_uniform(out_chan*5*5, 10)
+    if has_batchnorm:
+      self.bn1 = BatchNorm2d(inter_chan)
+      self.bn2 = BatchNorm2d(out_chan)
+    else:
+      self.bn1, self.bn2 = lambda x: x, lambda x: x
+
+  def parameters(self):
+    return get_parameters(self)
+
+  def forward(self, x:Tensor):
+    x = x.reshape(shape=(-1, 1, 28, 28)) # hacks
+    x = self.bn1(x.conv2d(self.c1)).relu().max_pool2d()
+    x = self.bn2(x.conv2d(self.c2)).relu().max_pool2d()
+    x = x.reshape(shape=[x.shape[0], -1])
+    return x.dot(self.l1)
+
+@slow
+class TestMNIST(unittest.TestCase):
+  def test_sgd_onestep(self):
+    np.random.seed(1337)
+    model = TinyBobNet()
+    optimizer = optim.SGD(model.parameters(), lr=0.001)
+    train(model, X_train, Y_train, optimizer, BS=69, steps=1)
+    for p in model.parameters(): p.realize()
+
+  def test_sgd_threestep(self):
+    np.random.seed(1337)
+    model = TinyBobNet()
+    optimizer = optim.SGD(model.parameters(), lr=0.001)
+    train(model, X_train, Y_train, optimizer, BS=69, steps=3)
+
+  def test_sgd_sixstep(self):
+    np.random.seed(1337)
+    model = TinyBobNet()
+    optimizer = optim.SGD(model.parameters(), lr=0.001)
+    train(model, X_train, Y_train, optimizer, BS=69, steps=6, noloss=True)
+
+  def test_adam_onestep(self):
+    np.random.seed(1337)
+    model = TinyBobNet()
+    optimizer = optim.Adam(model.parameters(), lr=0.001)
+    train(model, X_train, Y_train, optimizer, BS=69, steps=1)
+    for p in model.parameters(): p.realize()
+
+  def test_adam_threestep(self):
+    np.random.seed(1337)
+    model = TinyBobNet()
+    optimizer = optim.Adam(model.parameters(), lr=0.001)
+    train(model, X_train, Y_train, optimizer, BS=69, steps=3)
+
+  def test_conv_onestep(self):
+    np.random.seed(1337)
+    model = TinyConvNet()
+    optimizer = optim.SGD(model.parameters(), lr=0.001)
+    train(model, X_train, Y_train, optimizer, BS=69, steps=1, noloss=True)
+    for p in model.parameters(): p.realize()
+
+  def test_conv(self):
+    np.random.seed(1337)
+    model = TinyConvNet()
+    optimizer = optim.Adam(model.parameters(), lr=0.001)
+    train(model, X_train, Y_train, optimizer, steps=100)
+    assert evaluate(model, X_test, Y_test) > 0.93   # torch gets 0.9415 sometimes
+
+  def test_conv_with_bn(self):
+    np.random.seed(1337)
+    model = TinyConvNet(has_batchnorm=True)
+    optimizer = optim.AdamW(model.parameters(), lr=0.003)
+    train(model, X_train, Y_train, optimizer, steps=200)
+    assert evaluate(model, X_test, Y_test) > 0.94
+
+  def test_sgd(self):
+    np.random.seed(1337)
+    model = TinyBobNet()
+    optimizer = optim.SGD(model.parameters(), lr=0.001)
+    train(model, X_train, Y_train, optimizer, steps=600)
+    assert evaluate(model, X_test, Y_test) > 0.94   # CPU gets 0.9494 sometimes
+
+if __name__ == '__main__':
+  unittest.main()

artifacts/package_sources/tinygrad/test/models/test_onnx.py

@@ -0,0 +1,93 @@
+#!/usr/bin/env python
+import unittest
+import numpy as np
+from tinygrad.nn.onnx import OnnxRunner
+from tinygrad.device import Device
+from tinygrad.helpers import fetch, Context
+
+from extra.onnx_helpers import validate
+from extra.huggingface_onnx.huggingface_manager import DOWNLOADS_DIR, snapshot_download_with_retry
+
+def run_onnx_torch(onnx_model, inputs):
+  import torch
+  from onnx2torch import convert
+  torch_model = convert(onnx_model).float()
+  with torch.no_grad():
+    torch_out = torch_model(*[torch.tensor(x) for x in inputs.values()])
+  return torch_out
+
+np.random.seed(1337)
+
+class TestOnnxModel(unittest.TestCase):
+  @unittest.skip("slow")
+  def test_efficientnet(self):
+    input_name, input_new = "images:0", True
+    self._test_model(
+      fetch("https://github.com/onnx/models/raw/main/validated/vision/classification/efficientnet-lite4/model/efficientnet-lite4-11.onnx"),
+      input_name, input_new)
+
+  @unittest.skip("TODO: FIX THIS IT CAUSES SEGFAULT")
+  def test_shufflenet(self):
+    input_name, input_new = "gpu_0/data_0", False
+    self._test_model(
+      fetch("https://github.com/onnx/models/raw/main/validated/vision/classification/shufflenet/model/shufflenet-9.onnx"),
+      input_name, input_new)
+
+  @unittest.skip("test is very slow")
+  def test_resnet(self):
+    # NOTE: many onnx models can't be run right now due to max pool with strides != kernel_size
+    input_name, input_new = "data", False
+    self._test_model(
+      fetch("https://github.com/onnx/models/raw/main/validated/vision/classification/resnet/model/resnet18-v2-7.onnx"),
+      input_name, input_new)
+
+  def _test_model(self, fn, input_name, input_new, debug=False):
+    run_onnx = OnnxRunner(fn)
+    print("onnx loaded")
+    from test.models.test_efficientnet import chicken_img, car_img, preprocess, _LABELS
+
+    def run(img):
+      inputs = {input_name: preprocess(img, new=input_new)}
+      tinygrad_out = list(run_onnx(inputs, debug=debug).values())[0].numpy()
+      return tinygrad_out.argmax()
+
+    cls = run(chicken_img)
+    print(cls, _LABELS[cls])
+    assert _LABELS[cls] == "hen" or _LABELS[cls] == "cock"
+    cls = run(car_img)
+    print(cls, _LABELS[cls])
+    assert "car" in _LABELS[cls] or _LABELS[cls] == "convertible"
+
+@unittest.skipUnless(Device.DEFAULT == "METAL", "only run on METAL")
+class TestHuggingFaceOnnxModels(unittest.TestCase):
+  @classmethod
+  def setUpClass(cls):
+    cls._ctx = Context(MAX_BUFFER_SIZE=0)
+    cls._ctx.__enter__()
+
+  @classmethod
+  def tearDownClass(cls):
+    cls._ctx.__exit__()
+
+  def _validate(self, repo_id, model_file, custom_inputs, rtol=1e-4, atol=1e-4):
+    onnx_model_path = snapshot_download_with_retry(
+      repo_id=repo_id,
+      allow_patterns=["*.onnx", "*.onnx_data"],
+      local_dir=DOWNLOADS_DIR / repo_id
+    )
+    onnx_model_path = onnx_model_path / model_file
+    file_size = onnx_model_path.stat().st_size
+    print(f"Validating model: {repo_id}/{model_file} ({file_size/1e6:.2f}M)")
+    validate(onnx_model_path, custom_inputs, rtol=rtol, atol=atol)
+
+  def test_xlm_roberta_large(self):
+    repo_id = "FacebookAI/xlm-roberta-large"
+    model_file = "onnx/model.onnx"
+    custom_inputs = {
+      "input_ids": np.random.randint(0, 250002, (1, 11), dtype=np.int64),
+      "attention_mask": np.ones((1, 11), dtype=np.int64),
+    }
+    self._validate(repo_id, model_file, custom_inputs, atol=1e-3)
+
+if __name__ == "__main__":
+  unittest.main()

artifacts/package_sources/tinygrad/test/models/test_rnnt.py

@@ -0,0 +1,47 @@
+#!/usr/bin/env python
+import unittest
+from tinygrad import Tensor
+from extra.models.rnnt import LSTM
+import numpy as np
+import torch
+
+class TestRNNT(unittest.TestCase):
+  def test_lstm(self):
+    BS, SQ, IS, HS, L = 2, 20, 40, 128, 2
+
+    # create in torch
+    with torch.no_grad():
+      torch_layer = torch.nn.LSTM(IS, HS, L)
+
+    # create in tinygrad
+    layer = LSTM(IS, HS, L, 0.0)
+
+    # copy weights
+    with torch.no_grad():
+      layer.cells[0].weights_ih.assign(Tensor(torch_layer.weight_ih_l0.numpy()))
+      layer.cells[0].weights_hh.assign(Tensor(torch_layer.weight_hh_l0.numpy()))
+      layer.cells[0].bias_ih.assign(Tensor(torch_layer.bias_ih_l0.numpy()))
+      layer.cells[0].bias_hh.assign(Tensor(torch_layer.bias_hh_l0.numpy()))
+      layer.cells[1].weights_ih.assign(Tensor(torch_layer.weight_ih_l1.numpy()))
+      layer.cells[1].weights_hh.assign(Tensor(torch_layer.weight_hh_l1.numpy()))
+      layer.cells[1].bias_ih.assign(Tensor(torch_layer.bias_ih_l1.numpy()))
+      layer.cells[1].bias_hh.assign(Tensor(torch_layer.bias_hh_l1.numpy()))
+
+    # test initial hidden
+    for _ in range(3):
+      x = Tensor.randn(SQ, BS, IS)
+      z, hc = layer(x, None)
+      torch_x = torch.tensor(x.numpy())
+      torch_z, torch_hc = torch_layer(torch_x)
+      np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-3, rtol=5e-3)
+
+    # test passing hidden
+    for _ in range(3):
+      x = Tensor.randn(SQ, BS, IS)
+      z, hc = layer(x, hc)
+      torch_x = torch.tensor(x.numpy())
+      torch_z, torch_hc = torch_layer(torch_x, torch_hc)
+      np.testing.assert_allclose(z.numpy(), torch_z.detach().numpy(), atol=5e-3, rtol=5e-3)
+
+if __name__ == '__main__':
+  unittest.main()

artifacts/package_sources/tinygrad/test/models/test_train.py

@@ -0,0 +1,83 @@
+import unittest, time
+import numpy as np
+from tinygrad import Device
+from tinygrad.nn.state import get_parameters
+from tinygrad.nn import optim
+from tinygrad.helpers import getenv
+from test.helpers import slow
+from extra.training import train
+from extra.models.convnext import ConvNeXt
+from extra.models.efficientnet import EfficientNet
+from extra.models.transformer import Transformer
+from extra.models.vit import ViT
+from extra.models.resnet import ResNet18
+
+BS = getenv("BS", 2)
+
+def train_one_step(model,X,Y):
+  params = get_parameters(model)
+  pcount = 0
+  for p in params:
+    pcount += np.prod(p.shape)
+  optimizer = optim.SGD(params, lr=0.001)
+  print("stepping %r with %.1fM params bs %d" % (type(model), pcount/1e6, BS))
+  st = time.time()
+  train(model, X, Y, optimizer, steps=1, BS=BS)
+  et = time.time()-st
+  print("done in %.2f ms" % (et*1000.))
+
+def check_gc():
+  if Device.DEFAULT == "CL":
+    from extra.introspection import print_objects
+    assert print_objects() == 0
+
+class TestTrain(unittest.TestCase):
+  def test_convnext(self):
+    model = ConvNeXt(depths=[1], dims=[16])
+    X = np.zeros((BS,3,224,224), dtype=np.float32)
+    Y = np.zeros((BS), dtype=np.int32)
+    train_one_step(model,X,Y)
+    check_gc()
+
+  @slow
+  def test_efficientnet(self):
+    model = EfficientNet(0)
+    X = np.zeros((BS,3,224,224), dtype=np.float32)
+    Y = np.zeros((BS), dtype=np.int32)
+    train_one_step(model,X,Y)
+    check_gc()
+
+  @slow
+  def test_vit(self):
+    model = ViT()
+    X = np.zeros((BS,3,224,224), dtype=np.float32)
+    Y = np.zeros((BS,), dtype=np.int32)
+    train_one_step(model,X,Y)
+    check_gc()
+
+  @slow
+  def test_transformer(self):
+    # this should be small GPT-2, but the param count is wrong
+    # (real ff_dim is 768*4)
+    model = Transformer(syms=10, maxlen=6, layers=12, embed_dim=768, num_heads=12, ff_dim=768//4)
+    X = np.zeros((BS,6), dtype=np.float32)
+    Y = np.zeros((BS,6), dtype=np.int32)
+    train_one_step(model,X,Y)
+    check_gc()
+
+  @slow
+  def test_resnet(self):
+    X = np.zeros((BS, 3, 224, 224), dtype=np.float32)
+    Y = np.zeros((BS), dtype=np.int32)
+    for resnet_v in [ResNet18]:
+      model = resnet_v()
+      model.load_from_pretrained()
+      train_one_step(model, X, Y)
+    check_gc()
+
+  def test_bert(self):
+    # TODO: write this
+    pass
+
+if __name__ == '__main__':
+  unittest.main()

artifacts/package_sources/tinygrad/test/models/test_waifu2x.py

@@ -0,0 +1,24 @@
+#!/usr/bin/env python
+import pathlib
+import unittest
+import numpy as np
+from tinygrad.tensor import Tensor
+
+class TestVGG7(unittest.TestCase):
+  def test_vgg7(self):
+    from examples.vgg7_helpers.waifu2x import Vgg7, image_load
+
+    # Create in tinygrad
+    Tensor.manual_seed(1337)
+    mdl = Vgg7()
+    mdl.load_from_pretrained()
+
+    # Scale up an image
+    test_x = image_load(pathlib.Path(__file__).parent / 'waifu2x/input.png')
+    test_y = image_load(pathlib.Path(__file__).parent / 'waifu2x/output.png')
+    scaled = mdl.forward_tiled(test_x, 156)
+    scaled = np.fmax(0, np.fmin(1, scaled))
+    np.testing.assert_allclose(scaled, test_y, atol=5e-3, rtol=5e-3)
+
+if __name__ == '__main__':
+  unittest.main()

artifacts/package_sources/tinygrad/test/models/test_whisper.py

@@ -0,0 +1,153 @@
+import unittest
+import pathlib
+from examples.whisper import init_whisper, load_file_waveform, transcribe_file, transcribe_waveform
+from examples.audio_helpers import mel
+import examples.mlperf.metrics as metrics
+from tinygrad.helpers import fetch
+from test.helpers import slow
+from tinygrad import Tensor, Device, dtypes
+import numpy as np
+
+# Audio generated with the command on MacOS:
+# say "Could you please let me out of the box?" --file-format=WAVE  --data-format=LEUI8@16000 -o test
+# We use the WAVE type because it's easier to decode in CI test environments
+TEST_FILE_1 = str(pathlib.Path(__file__).parent / "whisper/test.wav")
+TRANSCRIPTION_1 = "Could you please let me out of the box?"
+TEST_FILE_2 = str(pathlib.Path(__file__).parent / "whisper/test2.wav")
+TRANSCRIPTION_2 = "a slightly longer audio file so that we can test batch transcriptions of varying length."
+# TODO this file will possibly not survive long. find another 1-2 minute sound file online to transcribe
+TEST_FILE_3_URL = 'https://homepage.ntu.edu.tw/~karchung/miniconversations/mc45.mp3'
+TRANSCRIPTION_3 = """Just lie back and relax.
+Is the level of pressure about right?
+Yes, it's fine. And I'd like conditioner, please.
+Sure. I'm going to start the second lathering now.
+Would you like some Q-tips?
+How'd you like it cut?
+I'd like my bangs and the back trimmed,
+and I'd like the rest thinned out a bit and layered.
+Where would you like the part?
+On the left, right about here.
+Here, have a look. What do you think?
+It's fine. Here's thousand NT dollars.
+It's 30 NT extra for the rinse. Here's your change and receipt.
+Thank you, and please come again!
+So, how do you like it?
+It could have been worse. But you'll notice that I didn't ask her for her card.
+Hmm, yeah.
+Mm, maybe you can try that place over there next time."""
+
+TRANSCRIPTION_3_ALT = "Just lie back and relax. Is the level of pressure about right? Yes, it's fine. And I'd like conditioner please. Sure. I'm going to start the second lathering now. Would you like some Q-tips? How'd you like it cut? I'd like my bangs on the back trimmed, and I'd like the rest to stand out a bit and layered. Where would you like the part? On the left, right about here. Here. Have a look. What do you think? It's fine. Here's a thousand and eighty dollars. It's thirty and t extra for the rants. Here's your change and receipt. Thank you, and please come again. So how do you like it? It could have been worse, but you'll notice that I didn't ask her for her card. Hmm, yeah. Maybe you can try that place over there next time." #noqa: E501
+# NOTE: same as TRANSCRIPTION_3 but with minor changes that should only amount to ~0.079 WER difference (see test_wer_same)
+# 'and'     --> 'on'
+# 'thinned' --> 'to stand'
+# 'nt'      --> 'and eighty'
+# '30 nt'   --> 'thirty and t'
+# 'rinse'   --> 'rants'
+# 'mm'      --> ''
+
+def wer_helper(result: str, reference: str)->float:
+  result = metrics.normalize_string(result)
+  reference = metrics.normalize_string(reference)
+  wer, _, _ = metrics.word_error_rate([result], [reference])
+  return wer
+
+# TODO: WEBGPU GPU dispatch dimensions limit
+@unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU GPU dispatch dimensions limit")
+class TestWhisper(unittest.TestCase):
+  @classmethod
+  def setUpClass(cls):
+    model, enc = init_whisper("tiny.en", batch_size=2)
+    cls.model = model
+    cls.enc = enc
+
+  @classmethod
+  def tearDownClass(cls):
+    del cls.model
+    del cls.enc
+
+  def assertWER(self, actual: str, expected: str, threshold: float):
+    __tracebackhide__ = True  # Hide traceback for py.test
+    wer = wer_helper(actual, expected)
+    if wer > threshold:
+      err = f"WER={wer:.3f} > {threshold}"
+      raise AssertionError(
+        err
+      )
+
+  def test_transcribe_file1(self):
+    self.assertEqual(transcribe_file(self.model, self.enc, TEST_FILE_1),  TRANSCRIPTION_1)
+
+  @slow
+  def test_transcribe_file2(self):
+    self.assertEqual(transcribe_file(self.model, self.enc, TEST_FILE_2),  TRANSCRIPTION_2)
+
+  @slow
+  def test_transcribe_batch12(self):
+    waveforms = [load_file_waveform(TEST_FILE_1), load_file_waveform(TEST_FILE_2)]
+    transcriptions = transcribe_waveform(self.model, self.enc, waveforms)
+    self.assertEqual(2, len(transcriptions))
+    self.assertEqual(TRANSCRIPTION_1,  transcriptions[0])
+    self.assertEqual(TRANSCRIPTION_2,  transcriptions[1])
+
+  def test_transcribe_batch21(self):
+    waveforms = [load_file_waveform(TEST_FILE_2), load_file_waveform(TEST_FILE_1)]
+    transcriptions = transcribe_waveform(self.model, self.enc, waveforms)
+    self.assertEqual(2, len(transcriptions))
+    self.assertEqual(TRANSCRIPTION_2,  transcriptions[0])
+    self.assertEqual(TRANSCRIPTION_1,  transcriptions[1])
+
+  @unittest.skip("file 3 url is broken")
+  @slow
+  def test_transcribe_long(self):
+    waveform = [load_file_waveform(fetch(TEST_FILE_3_URL))]
+    transcription = transcribe_waveform(self.model, self.enc, waveform)
+    self.assertWER(transcription, TRANSCRIPTION_3, 0.085)
+
+  @unittest.skip("file 3 url is broken")
+  @slow
+  def test_transcribe_long_no_batch(self):
+    waveforms = [load_file_waveform(fetch(TEST_FILE_3_URL)), load_file_waveform(TEST_FILE_1)]
+
+    trancriptions = transcribe_waveform(self.model, self.enc, waveforms)
+    self.assertEqual(2, len(trancriptions))
+    self.assertWER(trancriptions[0], TRANSCRIPTION_3, 0.085)
+    self.assertEqual(TRANSCRIPTION_1, trancriptions[1])
+
+  def test_wer_same(self):
+    reference = TRANSCRIPTION_3
+    self.assertWER(TRANSCRIPTION_3_ALT, reference, 0.079)
+
+  def test_wer_different(self):
+    reference = TRANSCRIPTION_3
+    self.assertWER("[no speech]", reference, 1.0)
+
+  def test_wer_different_2(self):
+    reference = TRANSCRIPTION_3
+    self.assertWER("", reference, 1.0)
+
+  def test_wer_different_3(self):
+    reference = TRANSCRIPTION_3
+    self.assertWER(reference[:len(reference)//2], reference, 0.524)
+
+  def test_mel_filters(self):
+    # reference = librosa.filters.mel(sr=16000, n_fft=16, n_mels=16)
+    reference = Tensor([[-0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                        [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                        [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                        [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                        [0.0, 0.0021111054811626673, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                        [0.0, 0.003133024089038372, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                        [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                        [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                        [0.0, 0.0, 0.0017568661132827401, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                        [0.0, 0.0, 0.0009823603322729468, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
+                        [0.0, 0.0, 0.0, 0.0007768510840833187, 0.0, 0.0, 0.0, 0.0, 0.0],
+                        [0.0, 0.0, 0.0, 0.0010490329004824162, 0.0, 0.0, 0.0, 0.0, 0.0],
+                        [0.0, 0.0, 0.0, 0.0, 0.0011341988574713469, 0.0, 0.0, 0.0, 0.0],
+                        [0.0, 0.0, 0.0, 0.0, 0.000231665835599415, 0.0006950111710466444, 0.0, 0.0, 0.0],
+                        [0.0, 0.0, 0.0, 0.0, 0.0, 0.00040073052514344454, 0.0005822855746373534, 0.0, 0.0],
+                        [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.00033081238507293165, 0.0006097797304391861, 0.0]])
+    np.testing.assert_allclose(mel(sr=16000, n_fft=16, n_mels=16, dtype=dtypes.float32).numpy(), reference.numpy(), atol=1e-6)
+
+if __name__ == '__main__':
+  unittest.main()