geohot/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Actions Activity

Increase op limit in test_index_mnist to pass on webgpu (#10504)

Browse Source 
* Increase op limit to enable  mnist indexing on webgpu

* Only relax op_limit on WebGPU
This commit is contained in:
Ahmed Harmouche
2025-05-24 09:37:31 -04:00
committed by GitHub
parent c472ab636c
commit bbb6deff53
1 changed files with 2 additions and 2 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/test_arange.py
+2 -2
View File
@@ -176,6 +176,8 @@ class TestIndexing(unittest.TestCase):
@unittest.skipIf(getenv("PTX"), "broken on ptx for some reason")
def test_index_mnist(self, noopt=1, op_limit=512*784*13, split_reduceop=0):
# WEBGPU generates more ops due to bitpacking of < 4-byte dtypes
if Device.DEFAULT == "WEBGPU": op_limit *= 15
from tinygrad.nn.datasets import mnist
X_train, Y_train, _, _ = mnist()
with Context(NOOPT=noopt, FUSE_ARANGE=1, SPLIT_REDUCEOP=split_reduceop):
@@ -190,9 +192,7 @@ class TestIndexing(unittest.TestCase):
# TODO: fix these on WEBGPU, it looks like it has to do with packed stuff
@unittest.skipIf(getenv("WEBGPU"), "broken on webgpu for some reason")
def test_index_mnist_opt(self): self.test_index_mnist(0)
@unittest.skipIf(getenv("WEBGPU"), "broken on webgpu for some reason")
def test_index_mnist_split(self): self.test_index_mnist(1, split_reduceop=1)
@unittest.skipIf(getenv("WEBGPU"), "broken on webgpu for some reason")
def test_index_mnist_opt_split(self): self.test_index_mnist(0, split_reduceop=1)
@unittest.skipIf(getenv("PTX"), "broken on ptx for some reason")