few minor code cleanups [pr] (#8267)

This commit is contained in:
chenyu
2024-12-15 23:44:51 -05:00
committed by GitHub
parent 9789a83064
commit f05fd118a2
3 changed files with 3 additions and 6 deletions
-2
View File
@@ -1,6 +1,5 @@
import unittest
from tinygrad import Variable
from tinygrad.helpers import getenv
from tinygrad.tensor import Tensor
from examples.gpt2 import Attention
import numpy as np
@@ -46,7 +45,6 @@ class TestSymbolicOps(unittest.TestCase):
expected = f(q, k, v).numpy()
np.testing.assert_allclose(symbolic, expected, atol=1e-6, rtol=1e-6)
@unittest.skipIf(getenv("MOCKHIP"), "MOCKHIP only compiles and does not run")
def test_attention_training(self):
with Tensor.train():
self.test_attention(dropout_p=0.0)
+1 -2
View File
@@ -1,9 +1,8 @@
import functools
import functools, struct
from tinygrad.device import Compiled, Allocator, Compiler
from tinygrad.renderer.wgsl import WGSLRenderer
from tinygrad.helpers import round_up
import wgpu
import struct
def create_uniform(wgpu_device, val) -> wgpu.GPUBuffer:
buf = wgpu_device.create_buffer(size=4, usage=wgpu.BufferUsage.UNIFORM | wgpu.BufferUsage.COPY_DST)
+2 -2
View File
@@ -43,7 +43,7 @@ def _reshape_mask(_mask:Optional[Tuple[Tuple[sint, sint], ...]], old_shape:Tuple
-> Optional[Tuple[Tuple[sint, sint], ...]]:
"""Returns the new mask if reshape is possible, and None if not possible."""
if _mask is None: return tuple((0, s) for s in new_shape)
if any(not isinstance(m[0], int) or not isinstance(m[1], int) for m in _mask): return None
if any(not all_int(m) for m in _mask): return None
if any(m[1] - m[0] < 1 for m in _mask): return ((0, 0),) * len(new_shape) # zero mask
new_mask: List[Tuple[int, int]] = []
@@ -332,9 +332,9 @@ class View:
acc = acc * new_dim
if not resolve(acc < real_size): new_stride = 0
if resolve(acc != merged_size): return None
new_strides = (0,) * (len(new_shape) - len(r_strides)) + tuple(r_strides[::-1])
if (new_mask:=_reshape_mask(self.mask, self.shape, new_shape)) is not None:
new_strides = (0,) * (len(new_shape) - len(r_strides)) + tuple(r_strides[::-1])
extra_offset = (sum(m[0] * s for m,s in zip(self.mask, self.strides)) if self.mask else 0) - \
(sum(m[0] * s for m,s in zip(new_mask, new_strides)))
return View.create(new_shape, new_strides, self.offset + extra_offset, new_mask)