mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-08-29 12:16:08 +00:00
few minor code cleanups [pr] (#8267)
This commit is contained in:
@@ -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,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)
|
||||
|
||||
@@ -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)
|
||||
|
||||
Reference in New Issue
Block a user