superrob/IQ.Pilot
1
0
Fork 0
forked from IQ.Lvbs/IQ.Pilot
Code Pull Requests Activity

IQ.Pilot Release Commit @ f2a861c

Browse Source
This commit is contained in:
IQ.Lvbs CI [bot]
2026-09-02 15:07:09 -05:00
parent b42569dbca
commit e8748fd704
5497 changed files with 316070 additions and 179848 deletions
Download Patch File Download Diff File Expand all files Collapse all files

0
artifacts/package_sources/tinygrad/test/unit/__init__.py Normal file
View File

85
artifacts/package_sources/tinygrad/test/unit/test_allreduce.py Normal file
View File

@@ -0,0 +1,85 @@
import unittest
from tinygrad import Tensor, UOp, dtypes
from tinygrad.helpers import Context
from tinygrad.uop.ops import Ops
class TestRingAllReduce(unittest.TestCase):
def test_schedule_ring(self):
with Context(RING=2):
N = 4
ds = tuple(f"CPU:{i}" for i in range(N))
t = Tensor.empty(N, N*100).shard(ds, axis=0).realize()
linear = t.sum(0).linear_with_vars()[0]
copies = [si for si in linear.src if si.src[0].op is Ops.COPY]
pairs = [(c.src[1].buffer.device, c.src[2].buffer.device) for c in copies]
# N*(N-1) scatter reduce, and N*(N-1) allgather
self.assertEqual(len(pairs), N*(N-1)*2)
# copy topology forms a ring
self.assertEqual(len(set(pairs)), N)
def test_schedule_all2all(self):
with Context(ALL2ALL=2):
N = 4
ds = tuple(f"CPU:{i}" for i in range(N))
t = Tensor.empty(N, N*100).shard(ds, axis=0).realize()
linear = t.sum(0).mul(2.0).contiguous().linear_with_vars()[0]
copies = [si for si in linear.src if si.src[0].op is Ops.COPY]
sinks = [si for si in linear.src if si.src[0].op is Ops.SINK]
self.assertEqual(len(copies), 24)
self.assertEqual(len(sinks), 26)
@Context(RING=0, ALL2ALL=0)
def test_schedule_naive(self):
N = 4
ds = tuple(f"NULL:{i}" for i in range(N))
t = Tensor.empty(N, 4096).shard(ds, axis=0).realize()
linear = t.sum(0).linear_with_vars()[0]
copies = [si for si in linear.src if si.src[0].op is Ops.COPY]
sinks = [si for si in linear.src if si.src[0].op is Ops.SINK]
pairs = [(c.src[1].buffer.device, c.src[2].buffer.device) for c in copies]
self.assertEqual(len(pairs), N*(N-1))
self.assertEqual(len(sinks), 2)
self.assertTrue(all(dst != src for dst, src in pairs))
def test_symbolic_shape(self):
rows = UOp.variable("rows", 1, 4).bind(3)
t = Tensor.ones(4, 4).shard(("CPU:0", "CPU:1"), axis=1).realize()
out = t[:rows].sum(1).realize()
self.assertEqual(out.shape, (rows,))
self.assertTrue((out == 4).all().item())
def test_correct_ring(self):
with Context(RING=2):
N = 4
ds = tuple(f"CPU:{i}" for i in range(N))
t = Tensor.ones(N, N*100).contiguous().shard(ds, axis=0).realize()
out = t.sum(0)
self.assertListEqual(out.tolist(), [4]*N*100)
class TestAllreduceCast(unittest.TestCase):
def _get_copy_dtypes(self, dtype, allreduce_cast):
ds = tuple(f"CPU:{i}" for i in range(2))
with Context(ALLREDUCE_CAST=allreduce_cast, RING=0, SCACHE=0):
t = Tensor.empty(4, 4, dtype=dtype).shard(ds, axis=0)
linear = t.sum(0).linear_with_vars()[0]
return {si.src[1].buffer.dtype.scalar() for si in linear.src if si.src[0].op is Ops.COPY}
def test_allreduce_cast_bf16(self):
# with ALLREDUCE_CAST, allreduce copies stay in bfloat16 instead of promoting to float32
self.assertNotIn(dtypes.float, self._get_copy_dtypes(dtypes.bfloat16, allreduce_cast=1))
self.assertIn(dtypes.float, self._get_copy_dtypes(dtypes.bfloat16, allreduce_cast=0))
def test_allreduce_cast_half(self):
self.assertNotIn(dtypes.float, self._get_copy_dtypes(dtypes.half, allreduce_cast=1))
self.assertIn(dtypes.float, self._get_copy_dtypes(dtypes.half, allreduce_cast=0))
def test_allreduce_cast_float32_noop(self):
# float32 should not be affected by ALLREDUCE_CAST (no promotion happens)
dtypes_on = self._get_copy_dtypes(dtypes.float, allreduce_cast=1)
dtypes_off = self._get_copy_dtypes(dtypes.float, allreduce_cast=0)
self.assertEqual(dtypes_on, dtypes_off)
if __name__ == '__main__':
unittest.main()

1079
artifacts/package_sources/tinygrad/test/unit/test_assign.py Normal file
View File

File diff suppressed because it is too large Load Diff

217
artifacts/package_sources/tinygrad/test/unit/test_attention.py Normal file
View File

@@ -0,0 +1,217 @@
import unittest
import numpy as np
from tinygrad import Tensor, dtypes
from tinygrad.llm.model import (
GatedDeltaNetBlock, SSMConfig, TransformerBlock, TransformerConfig,
apply_rope as apply_rope_new, precompute_freqs_cis, pairwise_topk,
)
def apply_rope(x:Tensor, start_pos:int):
B, H, T, Hd = x.shape
precompute_freqs_cis.cache_clear()
freqs_cis = precompute_freqs_cis(Hd, start_pos+T)[start_pos:start_pos+T]
return apply_rope_new(x, freqs_cis)
class TestAttention(unittest.TestCase):
def test_apply_rope(self):
x = Tensor.randn(1, 2, 4, 8, dtype=dtypes.float32)
result = apply_rope(x, 0)
self.assertEqual(result.shape, x.shape)
self.assertEqual(result.dtype, x.dtype)
self.assertGreater((result - apply_rope(x, 5)).abs().max().item(), 1e-6)
with self.assertRaises(AssertionError): apply_rope(Tensor.randn(1, 1, 4, 7, dtype=dtypes.float32), 0)
def test_partial_rope_in_attention(self):
dim, rope_dim, seqlen = 8, 4, 3
config = TransformerConfig(num_blocks=1, dim=dim, hidden_dim=16, n_heads=1, n_kv_heads=1,
norm_eps=1e-5, vocab_size=32, head_dim=dim, rope_theta=10000.0,
rope_dim=rope_dim, v_head_dim=dim, max_context=8)
block = TransformerBlock(config)
x = Tensor.randn(1, seqlen, dim, dtype=dtypes.float32)
x_norm = block.attn_norm(x)
k = block.attn_k(x_norm).reshape(1, seqlen, 1, dim).transpose(1, 2)
precompute_freqs_cis.cache_clear()
block.cache_kv = Tensor.empty(2, 1, 1, config.max_context, max(dim, config.v_head_dim), device=x.device)
block.freqs_cis = precompute_freqs_cis(rope_dim, config.max_context, config.rope_theta)
block._attention(x_norm, 0).realize()
expected = apply_rope_new(k[..., :rope_dim], block.freqs_cis[:seqlen]).cat(k[..., rope_dim:], dim=-1)
np.testing.assert_allclose(block.cache_kv[0, :, :, :seqlen, :].numpy(), expected.numpy(), rtol=1e-5, atol=1e-5)
class TestGatedDeltaNetBlock(unittest.TestCase):
def _tensor_linspace(self, start:float, stop:float, shape:tuple[int, ...]) -> Tensor:
return Tensor.linspace(start, stop, int(np.prod(shape)), dtype=dtypes.float32).reshape(*shape)
def _make_config(self, **kwargs):
return TransformerConfig(**({"num_blocks":1, "dim":4, "hidden_dim":8, "n_heads":1, "n_kv_heads":1,
"norm_eps":1e-5, "vocab_size":32, "head_dim":4, "rope_theta":10000.0,
"rope_dim":4, "v_head_dim":4, "max_context":4, "ssm_layers":(True,),
"ssm":SSMConfig(conv_kernel=2, state_size=2, group_count=1, time_step_rank=1, inner_size=2)} | kwargs))
def _make_block(self, config:TransformerConfig) -> GatedDeltaNetBlock:
block = GatedDeltaNetBlock(config, config.ssm)
block.attn_norm.weight = self._tensor_linspace(0.8, 1.2, (config.dim,))
block.attn_qkv.weight = self._tensor_linspace(-0.15, 0.2, (block.conv_channels, config.dim))
block.attn_gate.weight = self._tensor_linspace(-0.1, 0.15, (config.ssm.inner_size, config.dim))
block.ssm_alpha.weight = self._tensor_linspace(-0.08, 0.12, (block.num_v_heads, config.dim))
block.ssm_beta.weight = self._tensor_linspace(-0.12, 0.07, (block.num_v_heads, config.dim))
block.ssm_conv1d["weight"] = self._tensor_linspace(-0.05, 0.05, (block.conv_channels, block.ssm_conv_kernel))
block.ssm_dt["bias"] = self._tensor_linspace(-0.1, 0.1, (block.num_v_heads,))
block.ssm_a = self._tensor_linspace(-0.1, -0.05, (block.num_v_heads,))
block.ssm_norm.weight = self._tensor_linspace(0.9, 1.1, (block.head_v_dim,))
block.ssm_out.weight = self._tensor_linspace(-0.2, 0.18, (config.dim, config.ssm.inner_size))
return block
def _run_attention(self, block:GatedDeltaNetBlock, x:Tensor, start_pos:int):
x_norm = block.attn_norm(x)
block._init_state(x_norm)
return block._attention(x_norm, start_pos).realize().numpy()
def _cache_views(self, block:GatedDeltaNetBlock) -> tuple[np.ndarray, np.ndarray]:
if hasattr(block, 'conv_state'):
return block.conv_state.numpy(), block.recurrent_state.numpy()
else:
conv_flat = (block.ssm_conv_kernel - 1) * block.conv_channels
cache = block.delta_cache.numpy()
conv_state = cache[:, :conv_flat].reshape(cache.shape[0], block.ssm_conv_kernel - 1, block.conv_channels)
recurrent_state = cache[:, conv_flat:].reshape(cache.shape[0], block.num_v_heads, block.head_v_dim, block.head_v_dim)
return conv_state, recurrent_state
def _linear_np(self, x:np.ndarray, weight:np.ndarray) -> np.ndarray:
return x.astype(np.float32) @ weight.T.astype(np.float32)
def _rms_norm_np(self, x:np.ndarray, weight:np.ndarray, eps:float) -> np.ndarray:
x_float = x.astype(np.float32)
return (x_float / np.sqrt((x_float * x_float).mean(axis=-1, keepdims=True) + eps)) * weight.astype(np.float32)
def _normalize_np(self, x:np.ndarray, eps:float=1e-12) -> np.ndarray:
return x / np.maximum(np.sqrt((x * x).sum(axis=-1, keepdims=True)), eps)
def _softplus_np(self, x:np.ndarray) -> np.ndarray:
return np.log1p(np.exp(-np.abs(x))) + np.maximum(x, 0)
def _silu_np(self, x:np.ndarray) -> np.ndarray:
return x / (1.0 + np.exp(-x))
def _naive_attention(self, block:GatedDeltaNetBlock, x:Tensor):
x_np = x.numpy().astype(np.float32)
B, T, _ = x_np.shape
conv_state = np.zeros((B, block.ssm_conv_kernel - 1, block.conv_channels), dtype=np.float32)
recurrent_state = np.zeros((B, block.num_v_heads, block.head_v_dim, block.head_v_dim), dtype=np.float32)
conv_weight = block.ssm_conv1d["weight"].numpy().astype(np.float32).T[None, :, :]
qkv_weight = block.attn_qkv.weight.numpy().astype(np.float32)
gate_weight = block.attn_gate.weight.numpy().astype(np.float32)
alpha_weight = block.ssm_alpha.weight.numpy().astype(np.float32)
beta_weight = block.ssm_beta.weight.numpy().astype(np.float32)
out_weight = block.ssm_out.weight.numpy().astype(np.float32)
dt_bias = block.ssm_dt["bias"].numpy().astype(np.float32)
ssm_a = block.ssm_a.numpy().astype(np.float32)
attn_norm_weight = block.attn_norm.weight.numpy().astype(np.float32)
ssm_norm_weight = block.ssm_norm.weight.numpy().astype(np.float32)
outputs, conv_states, recurrent_states = [], [], []
for t in range(T):
x_norm = self._rms_norm_np(x_np[:, t:t+1, :], attn_norm_weight, block.attn_norm.eps)
x_half = x_norm.astype(np.float16)
out_gate = self._linear_np(x_half, gate_weight).reshape(B, 1, block.num_v_heads, block.head_v_dim)
beta = 1.0 / (1.0 + np.exp(-self._linear_np(x_half, beta_weight))).reshape(B, block.num_v_heads, 1, 1)
alpha = np.exp((self._softplus_np(self._linear_np(x_half, alpha_weight) + dt_bias)).reshape(B, block.num_v_heads, 1, 1) *
ssm_a.reshape(1, block.num_v_heads, 1, 1))
conv_window = np.concatenate([conv_state, self._linear_np(x_half, qkv_weight)], axis=1)
conv_out = self._silu_np((conv_window * conv_weight).sum(axis=1))
q, k, v = np.split(conv_out, [block.q_dim, 2 * block.q_dim], axis=-1)
q = self._normalize_np(q.reshape(B, block.num_k_heads, block.head_k_dim))
k = self._normalize_np(k.reshape(B, block.num_k_heads, block.head_k_dim))
v = v.reshape(B, block.num_v_heads, block.head_v_dim)
if block.num_v_heads != block.num_k_heads:
k_repeat = block.num_v_heads // block.num_k_heads
q = np.repeat(q[:, None, :, :], k_repeat, axis=1).reshape(B, block.num_v_heads, block.head_k_dim)
k = np.repeat(k[:, None, :, :], k_repeat, axis=1).reshape(B, block.num_v_heads, block.head_k_dim)
q, k, v = (q * (block.head_k_dim ** -0.5))[..., None], k[..., None], v[..., None]
recurrent_state = recurrent_state * alpha
recurrent_state = recurrent_state + np.matmul((v - np.matmul(recurrent_state, k)) * beta, np.swapaxes(k, -1, -2))
core_attn_out = np.matmul(recurrent_state, q).squeeze(-1).reshape(B, 1, block.num_v_heads, block.head_v_dim)
core_attn_out = self._rms_norm_np(core_attn_out, ssm_norm_weight, block.ssm_norm.eps)
out = self._linear_np((core_attn_out * self._silu_np(out_gate)).reshape(B, 1, -1).astype(np.float16), out_weight)
conv_state = conv_window[:, 1:, :]
outputs.append(out)
conv_states.append(conv_state.copy())
recurrent_states.append(recurrent_state.copy())
return outputs, conv_states, recurrent_states
def test_gatedeltanet_reference_and_reset(self):
config = self._make_config(max_context=3)
block = self._make_block(config)
x = Tensor.linspace(-1.0, 1.0, 3 * config.dim, dtype=dtypes.float32).reshape(1, 3, config.dim)
expected_outs, expected_conv, expected_recurrent = self._naive_attention(block, x)
for step in range(x.shape[1]):
out = self._run_attention(block, x[:, step:step+1], step)
conv_state, recurrent_state = self._cache_views(block)
np.testing.assert_allclose(out, expected_outs[step], rtol=1e-3, atol=1e-3,
err_msg=f"GatedDeltaNet output mismatch at step {step}")
np.testing.assert_allclose(conv_state, expected_conv[step], rtol=1e-3, atol=1e-3,
err_msg=f"GatedDeltaNet conv cache mismatch at step {step}")
np.testing.assert_allclose(recurrent_state, expected_recurrent[step], rtol=1e-3, atol=1e-3,
err_msg=f"GatedDeltaNet recurrent cache mismatch at step {step}")
warmup = Tensor.linspace(-0.5, 0.5, 2 * config.dim, dtype=dtypes.float32).reshape(1, 2, config.dim)
prompt = Tensor.linspace(0.75, -0.75, 2 * config.dim, dtype=dtypes.float32).reshape(1, 2, config.dim)
for i in range(warmup.shape[1]): self._run_attention(block, warmup[:, i:i+1], i)
Tensor.realize(*block._state_reset_ops())
expected_outs, expected_conv, expected_recurrent = self._naive_attention(block, prompt)
for step in range(prompt.shape[1]):
out = self._run_attention(block, prompt[:, step:step+1], step)
conv_state, recurrent_state = self._cache_views(block)
np.testing.assert_allclose(out, expected_outs[step], rtol=1e-3, atol=1e-3,
err_msg=f"GatedDeltaNet reset output mismatch at step {step}")
np.testing.assert_allclose(conv_state, expected_conv[step], rtol=1e-3, atol=1e-3,
err_msg=f"GatedDeltaNet reset conv cache mismatch at step {step}")
np.testing.assert_allclose(recurrent_state, expected_recurrent[step], rtol=1e-3, atol=1e-3,
err_msg=f"GatedDeltaNet reset recurrent cache mismatch at step {step}")
def test_kda_channel_decay(self):
config = self._make_config(n_heads=2, ssm=SSMConfig(conv_kernel=2, state_size=2, group_count=2, time_step_rank=2, inner_size=4, kda=True))
block, x = GatedDeltaNetBlock(config, config.ssm), Tensor([[[1., 2., 0., 0.]]])
# f_b(f_a(x)) = [1, 2, 3, 4]
block.ssm_f_a.weight = Tensor([[1., 0., 0., 0.], [0., 1., 0., 0.]])
block.ssm_f_b.weight = Tensor([[1., 0.], [0., 1.], [1., 1.], [2., 1.]])
block._init_state(x)
initial_state = Tensor.arange(8, dtype=dtypes.float32).reshape(1, 2, 2, 2)
block.recurrent_state.assign(initial_state).realize()
block.ssm_a = Tensor([[-1.], [-1.]])
block._attention(x, 0).realize()
alpha = np.exp(-self._softplus_np(np.arange(1, 5)).reshape(1, 2, 1, 2))
np.testing.assert_allclose(block.recurrent_state.numpy(), initial_state.numpy() * alpha, rtol=1e-5, atol=1e-5)
class TestPairwiseTopk(unittest.TestCase):
def test_basic_topk(self):
x = Tensor([[[1.0, 3.0, 2.0, 5.0, 4.0]]])
vals, sel = pairwise_topk(x, 3)
np.testing.assert_allclose(vals.numpy(), [[[3.0, 4.0, 5.0]]])
np.testing.assert_equal(sel.numpy(), [[[1, 4, 3]]])
def test_duplicates(self):
x = Tensor([[[5.0, 5.0, 3.0, 5.0]]])
vals, sel = pairwise_topk(x, 2)
np.testing.assert_allclose(vals.numpy(), [[[5.0, 5.0]]])
np.testing.assert_equal(sel.numpy(), [[[1, 0]]])
def test_matches_numpy(self):
np.random.seed(42)
data = np.random.randn(4, 2, 16).astype(np.float32)
vals, sel = pairwise_topk(Tensor(data), 5)
for b in range(4):
for t in range(2):
expected = set(np.argsort(-data[b, t])[:5].tolist())
self.assertEqual(set(sel.numpy()[b, t].tolist()), expected)
np.testing.assert_allclose(vals.numpy()[b, t], data[b, t][sel.numpy()[b, t]])
if __name__ == '__main__':
unittest.main()

363
artifacts/package_sources/tinygrad/test/unit/test_call.py Normal file
View File

@@ -0,0 +1,363 @@
import unittest
import numpy as np
from tinygrad import Tensor, function
from tinygrad.dtype import dtypes
from tinygrad.uop.ops import UOp, Ops
class TestCall(unittest.TestCase):
def test_call_plus(self):
a = Tensor.randn(10, 10)
b = Tensor.randn(10, 10)
Tensor.realize(a,b)
# we define a plus function
plus_fxn = UOp.param(0, dtypes.float, (10,10)) + UOp.param(1, dtypes.float, (10,10))
c = Tensor.call(a, b, fxn=plus_fxn)
np.testing.assert_equal(c.numpy(), (a+b).numpy())
def test_call_plus_backward(self):
a = Tensor.ones(10, 10)
b = Tensor.ones(10, 10)
(a+b).mean().backward()
gt_a_grad = a.grad.numpy()
gt_b_grad = b.grad.numpy()
a.grad, b.grad = None, None
# this is the gradient for +
def grad_fxn(grad:UOp, call:UOp): return (grad, grad)
# we define a plus function
plus_fxn = UOp.param(0, dtypes.float, (10,10)) + UOp.param(1, dtypes.float, (10,10))
c = Tensor.call(a, b, fxn=plus_fxn, grad_fxn=grad_fxn)
c.mean().backward()
np.testing.assert_allclose(a.grad.numpy(), gt_a_grad, rtol=1e-5)
np.testing.assert_allclose(b.grad.numpy(), gt_b_grad, rtol=1e-5)
def test_call_plus_backward_auto(self):
a = Tensor.ones(10, 10)
b = Tensor.ones(10, 10)
(a+b).mean().backward()
gt_a_grad = a.grad.numpy()
gt_b_grad = b.grad.numpy()
a.grad, b.grad = None, None
plus_fxn = UOp.param(0, dtypes.float, (10,10)) + UOp.param(1, dtypes.float, (10,10))
c = Tensor.call(a, b, fxn=plus_fxn)
c.mean().backward()
np.testing.assert_allclose(a.grad.numpy(), gt_a_grad, rtol=1e-5)
np.testing.assert_allclose(b.grad.numpy(), gt_b_grad, rtol=1e-5)
def test_call_gemm(self):
M, K, N = 4, 8, 4
a = Tensor.randn(M, K)
b = Tensor.randn(K, N)
Tensor.realize(a, b)
c = Tensor.call(a, b, fxn=a.as_param(0) @ b.as_param(1))
np.testing.assert_allclose(c.numpy(), a.numpy() @ b.numpy(), rtol=1e-5, atol=1e-6)
def test_call_gemm_uop(self):
M, K, N = 4, 8, 4
a = Tensor.randn(M, K)
b = Tensor.randn(K, N)
Tensor.realize(a, b)
# we define a gemm function
x = UOp.param(0, dtypes.float, shape=(M, K))
y = UOp.param(1, dtypes.float, shape=(K, N))
c = Tensor.call(a, b, fxn=x@y)
np.testing.assert_allclose(c.numpy(), a.numpy() @ b.numpy(), rtol=1e-5, atol=1e-6)
def test_call_complex_backward_auto(self):
# complex chain: (a*b + a).exp2() * b.reciprocal() - tests mul, add, exp2, reciprocal, param reuse
a = Tensor.randn(10, 10)
b = Tensor.randn(10, 10) + 2 # avoid div by zero
Tensor.realize(a, b)
((a*b + a).exp2() * b.reciprocal()).mean().backward()
gt_a_grad, gt_b_grad = a.grad.numpy(), b.grad.numpy()
a.grad, b.grad = None, None
p0, p1 = UOp.param(0, dtypes.float, (10,10)), UOp.param(1, dtypes.float, (10,10))
complex_fxn = (p0*p1 + p0).exp2() * p1.reciprocal()
c = Tensor.call(a, b, fxn=complex_fxn)
c.mean().backward()
np.testing.assert_allclose(a.grad.numpy(), gt_a_grad, rtol=1e-5)
np.testing.assert_allclose(b.grad.numpy(), gt_b_grad, rtol=1e-5)
def test_call_plus_sharded(self):
devs = ("CPU:0", "CPU:1")
a = Tensor.ones(10, 10).shard(devs, axis=0)
b = Tensor.ones(10, 10).shard(devs, axis=0)
Tensor.realize(a, b)
c = Tensor.call(a, b, fxn=a.as_param(0) + b.as_param(1))
np.testing.assert_equal(c.numpy(), 2 * np.ones((10, 10)))
class TestCallShape(unittest.TestCase):
def test_call_shape_int(self):
# fixed-shape function: shape passes through unchanged
@function
def f(x:Tensor) -> Tensor: return x * 2
self.assertEqual(f(Tensor.empty(4, 8)).shape, (4, 8))
def test_call_shape_param_substitution(self):
# symbolic shape dimension is substituted: inner PARAM replaced with the BIND arg
@function
def f(x:Tensor) -> Tensor: return x * 2
sz = UOp.variable("sz", 1, 8)
shape = f(Tensor.empty(8)[:sz.bind(5)]).shape
# the PARAM should be gone, replaced with the BIND from the call arg
self.assertIsInstance(shape[0], UOp)
self.assertNotEqual(shape[0].op, Ops.PARAM)
self.assertEqual(shape[0], sz.bind(5))
def test_call_shape_expr_substitution(self):
# expression containing PARAMs in shape gets fully substituted
@function
def f(x:Tensor) -> Tensor: return x + 1
sz = UOp.variable("sz", 1, 10)
shape = f(Tensor.empty(10, 4)[:sz.bind(3)]).shape
self.assertIsInstance(shape[0], UOp)
self.assertNotEqual(shape[0].op, Ops.PARAM)
self.assertEqual(shape[1], 4)
def test_call_shape_no_param_passthrough(self):
# a non-PARAM UOp shape element passes through unchanged
@function
def f(x:Tensor) -> Tensor: return x * 3
sz = UOp.variable("sz", 1, 8)
shape = f(Tensor.empty(8)[:sz.bind(5)]).shape
self.assertEqual(shape[0], sz.bind(5))
class TestCallSchedule(unittest.TestCase):
def test_reshape_precompile(self):
a = Tensor.empty(4, 8).realize()
a = a.reshape(4,4,2).assign(Tensor.empty(4,4,2)).reshape(8,4)
@function(precompile=True)
def s(x): return x.sum(axis=0)
(s(a)*3).realize()
def test_call_precompiled(self):
a = Tensor.empty(4, 8)
@function(precompile=True)
def s(x): return x*2
(s(a)*3).realize()
def test_double_call(self):
a = Tensor.empty(4, 8)
@function(precompile=True)
def s(x): return x*2
s(s(a)).realize()
def test_double_call_contiguous(self):
a = Tensor.empty(4, 8)
@function(precompile=True)
def s(x): return x*2
s(s(a).contiguous()).realize()
def test_call_double_gemm(self):
a = Tensor.randn(4, 8)
b = Tensor.randn(8, 12)
c = Tensor.randn(12, 16)
ref = Tensor.randn(4, 16)
Tensor.realize(a,b,c,ref)
@function(precompile=True)
def gemm(a:Tensor, b:Tensor, c:Tensor) -> Tensor: return (a@b)@c
out = gemm(a,b,c)
(out-ref).square().mean().backward()
out.realize(a.grad, b.grad, c.grad)
def test_precompile_symbolic_shape(self):
"""precompile with a symbolic-shaped input produces correct values and shape"""
@function(precompile=True)
def f(x:Tensor) -> Tensor: return x * 2
sz = UOp.variable("sz", 1, 8)
a = Tensor([1., 2., 3., 4., 5., 6., 7., 8.])[:sz.bind(5)]
out = f(a)
self.assertIsInstance(out.shape[0], UOp)
np.testing.assert_allclose(out[:5].numpy(), [2., 4., 6., 8., 10.])
def test_precompile_symbolic_shape_contiguous(self):
"""precompile with a .contiguous() inside the function body on a symbolic-shaped input"""
@function(precompile=True)
def f(x:Tensor) -> Tensor: return (x * 2).contiguous() + 1
sz = UOp.variable("sz", 1, 8)
a = Tensor([1., 2., 3., 4., 5., 6., 7., 8.])[:sz.bind(3)]
out = f(a)
self.assertIsInstance(out.shape[0], UOp)
np.testing.assert_allclose(out[:3].numpy(), [3., 5., 7.])
def test_precompile_symbolic_shape_chain(self):
"""precompiled symbolic result used in downstream ops (tests AFTER has correct symbolic shape)"""
@function(precompile=True)
def f(x:Tensor) -> Tensor: return x * 2
sz = UOp.variable("sz", 1, 8)
a = Tensor([1., 2., 3., 4., 5., 6., 7., 8.])[:sz.bind(4)]
out = f(a) + 10 # downstream op on the precompiled result
self.assertIsInstance(out.shape[0], UOp)
np.testing.assert_allclose(out[:4].numpy(), [12., 14., 16., 18.])
def test_precompile_bind_arg(self):
"""precompile with a BIND (scalar variable) as a function argument"""
@function(precompile=True)
def f(x:Tensor, scale:UOp) -> Tensor: return x * scale
v = UOp.variable("scale", 1, 100)
a = Tensor([1., 2., 3.])
out = f(a, v.bind(5))
np.testing.assert_allclose(out.numpy(), [5., 10., 15.])
def test_precompile_scoped_bind_arg(self):
@function(precompile=True)
def f(x:Tensor, scale:UOp) -> Tensor: return x * scale
a = Tensor.ones(3)
x = f(a, UOp.variable("scale_a", 1, 100).bind(2))
y = f(a, UOp.variable("scale_b", 1, 100).bind(3))
fx = next(u for u in x.uop.toposort() if u.op is Ops.FUNCTION)
fy = next(u for u in y.uop.toposort() if u.op is Ops.FUNCTION)
self.assertEqual(fx.src[0].key, fy.src[0].key)
np.testing.assert_equal(x.numpy(), [2, 2, 2])
np.testing.assert_equal(y.numpy(), [3, 3, 3])
def test_precompile_schedule_cache_hit(self):
"""two instances of the same @function should produce identical function body keys (schedule cache hit)"""
@function(precompile=True)
def f(x:Tensor) -> Tensor: return x + Tensor.full(x.shape, -1.0)
a = Tensor.empty(4, 8)
b = Tensor.empty(4, 8)
r0, r1 = f(a), f(b)
# find the FUNCTION nodes
c0 = next(u for u in r0.uop.toposort() if u.op is Ops.FUNCTION)
c1 = next(u for u in r1.uop.toposort() if u.op is Ops.FUNCTION)
# the function bodies (src[0]) should have identical keys
self.assertEqual(c0.src[0].key, c1.src[0].key)
def test_precompile_symbolic_2d(self):
"""precompile with symbolic shapes in 2D (tests debuf reshape with symbolic PARAM)"""
@function(precompile=True)
def f(x:Tensor) -> Tensor: return x * 2 + 1
sz = UOp.variable("sz", 1, 16)
a = Tensor.arange(16*4).reshape(16, 4).float().clone()[:sz.bind(5)]
out = f(a)
# result shape should have the symbolic dim, not the max
self.assertIsInstance(out.shape[0], UOp)
np.testing.assert_allclose(out[:5].numpy(), (np.arange(16*4).reshape(16, 4)[:5] * 2 + 1).astype(np.float32))
def test_precompile_multi_sharded(self):
@function(precompile=True)
def f(x:Tensor) -> Tensor: return x + 1
devs = ("CPU:0", "CPU:1")
a = Tensor.arange(8).reshape(4, 2).float().clone().shard(devs, axis=0)
out = f(a) + 2
np.testing.assert_allclose(out.numpy(), np.arange(8, dtype=np.float32).reshape(4, 2) + 3)
class TestCallMultiSharded(unittest.TestCase):
# TODO: multi-output + sharded needs per-device CALL execution, which requires reworking how MULTI propagates through TUPLE bodies
def test_tuple_sharded(self):
"""multi-output function with sharded input"""
devs = ("CPU:0", "CPU:1")
@function
def f(x:Tensor): return (x + 1, x * 2)
a = Tensor.arange(8).reshape(4, 2).float().clone().shard(devs, axis=0)
t1, t2 = f(a)
ref = np.arange(8, dtype=np.float32).reshape(4, 2)
np.testing.assert_allclose(t1.numpy(), ref + 1)
np.testing.assert_allclose(t2.numpy(), ref * 2)
def test_tuple_sharded_precompile(self):
"""multi-output precompiled function with sharded input"""
devs = ("CPU:0", "CPU:1")
@function(precompile=True)
def f(x:Tensor): return (x + 1, x * 2)
a = Tensor.arange(8).reshape(4, 2).float().clone().shard(devs, axis=0)
t1, t2 = f(a)
ref = np.arange(8, dtype=np.float32).reshape(4, 2)
np.testing.assert_allclose(t1.numpy(), ref + 1)
np.testing.assert_allclose(t2.numpy(), ref * 2)
def test_tuple_sharded_different_axis(self):
"""multi-output function where outputs have different sharding: one reduces on sharded axis, one doesn't"""
devs = ("CPU:0", "CPU:1")
@function
def f(x:Tensor): return (x.sum(axis=0), x.sum(axis=1))
a = Tensor.arange(8).reshape(4, 2).float().clone().shard(devs, axis=0)
t1, t2 = f(a)
ref = np.arange(8, dtype=np.float32).reshape(4, 2)
np.testing.assert_allclose(t1.numpy(), ref.sum(axis=0))
np.testing.assert_allclose(t2.numpy(), ref.sum(axis=1))
def test_tuple_sharded_different_ops(self):
"""multi-output function with different operations per output"""
devs = ("CPU:0", "CPU:1")
@function
def f(x:Tensor, y:Tensor): return (x + y, x * y)
a = Tensor.arange(8).reshape(4, 2).float().clone().shard(devs, axis=0)
b = Tensor.arange(8).reshape(4, 2).float().clone().shard(devs, axis=0) + 1
t1, t2 = f(a, b)
ref_a = np.arange(8, dtype=np.float32).reshape(4, 2)
ref_b = ref_a + 1
np.testing.assert_allclose(t1.numpy(), ref_a + ref_b)
np.testing.assert_allclose(t2.numpy(), ref_a * ref_b)
def test_tuple_sharded_mixed_use(self):
"""multi-output sharded results used in further computation"""
devs = ("CPU:0", "CPU:1")
@function
def f(x:Tensor): return (x + 1, x * 2)
a = Tensor.arange(8).reshape(4, 2).float().clone().shard(devs, axis=0)
t1, t2 = f(a)
out = (t1 + t2).sum()
ref = np.arange(8, dtype=np.float32).reshape(4, 2)
np.testing.assert_allclose(out.numpy(), ((ref + 1) + (ref * 2)).sum())
def test_tuple_sharded_outputs_different_axis(self):
"""multi-output function where the two outputs are sharded on different axes"""
devs = ("CPU:0", "CPU:1")
@function
def f(x:Tensor, y:Tensor): return (x + 1, y + 2)
a = Tensor.arange(8).reshape(4, 2).float().clone().shard(devs, axis=0)
b = Tensor.arange(8).reshape(4, 2).float().clone().shard(devs, axis=1)
t1, t2 = f(a, b)
ref_a = np.arange(8, dtype=np.float32).reshape(4, 2)
ref_b = np.arange(8, dtype=np.float32).reshape(4, 2)
np.testing.assert_allclose(t1.numpy(), ref_a + 1)
np.testing.assert_allclose(t2.numpy(), ref_b + 2)
def test_call_reduce_sharded(self):
devs = ("CPU:0", "CPU:1")
a = Tensor.ones(10, 10).shard(devs, axis=0)
Tensor.realize(a)
c = Tensor.call(a, fxn=a.as_param(0).sum(axis=0))
np.testing.assert_equal(c.numpy(), 10 * np.ones(10))
def test_call_reduce_sharded_mixed_args(self):
devs = ("CPU:0", "CPU:1")
a = Tensor.ones(10, 10).shard(devs, axis=0)
b = Tensor.ones(10).shard(devs, axis=None)
Tensor.realize(a, b)
c = Tensor.call(a, b, fxn=a.as_param(0).sum(axis=0) + b.as_param(1))
np.testing.assert_equal(c.numpy(), 11 * np.ones(10))
def test_call_reduce_sharded_backward(self):
devs = ("CPU:0", "CPU:1")
a = Tensor.randn(10, 10).shard(devs, axis=0)
b = Tensor.randn(10, 10).shard(devs, axis=0)
Tensor.realize(a, b)
def grad_fxn(grad, call):
a_arg, b_arg = call.src[1], call.src[2]
return (grad.expand(a_arg.shape) * b_arg, grad.expand(b_arg.shape) * a_arg)
body = (a.as_param(0) * b.as_param(1)).sum(axis=0)
c = Tensor.call(a, b, fxn=body, grad_fxn=grad_fxn)
c.sum().backward()
np.testing.assert_allclose(a.grad.numpy(), b.numpy(), rtol=1e-5)
np.testing.assert_allclose(b.grad.numpy(), a.numpy(), rtol=1e-5)
if __name__ == '__main__':
unittest.main()

111
artifacts/package_sources/tinygrad/test/unit/test_callify.py Normal file
View File

@@ -0,0 +1,111 @@
import unittest
from tinygrad import Tensor, dtypes
class TestCallify(unittest.TestCase):
def test_basic(self):
a = Tensor([1.,2,3])
b = Tensor([4.,5,6])
out = a + b
out.callify()
self.assertListEqual(out.tolist(), [5.0, 7.0, 9.0])
def test_const(self):
out = Tensor(2.0) + Tensor(3.0)
out.callify()
self.assertEqual(out.item(), 5.0)
def test_sum(self):
out = Tensor.ones(16).contiguous().sum()
out.callify()
self.assertEqual(out.item(), 16.0)
def test_multi_output(self):
a = Tensor([1.,2,3])
b = Tensor([4.,5,6])
c = a + b
d = a * b
c.callify(d)
self.assertListEqual(c.tolist(), [5.0, 7.0, 9.0])
self.assertListEqual(d.tolist(), [4.0, 10.0, 18.0])
def test_two_callify_independent(self):
a = Tensor([1.,2,3])
b = Tensor([4.,5,6])
c = a + b
c.callify()
d = Tensor([10.,20,30])
e = Tensor([1.,1,1])
f = d - e
f.callify()
self.assertListEqual(c.tolist(), [5.0, 7.0, 9.0])
self.assertListEqual(f.tolist(), [9.0, 19.0, 29.0])
def test_two_callify_shared_input(self):
a = Tensor([1.,2,3]).contiguous().realize()
b = a + 1
b.callify()
c = a * 2
c.callify()
self.assertListEqual(b.tolist(), [2.0, 3.0, 4.0])
self.assertListEqual(c.tolist(), [2.0, 4.0, 6.0])
def test_chained_callify(self):
a = Tensor([1.,2,3])
b = a + 1
b.callify()
b.realize()
c = b + 1
c.callify()
self.assertListEqual(c.tolist(), [3.0, 4.0, 5.0])
def test_gemm(self):
a = Tensor.ones(8, 8).contiguous()
b = Tensor.eye(8).contiguous()
out = a @ b
out.callify()
lst = out.tolist()
for y in range(8):
for x in range(8):
self.assertEqual(lst[y][x], 1.0)
def test_int_dtype(self):
a = Tensor([1,2,3], dtype=dtypes.int)
b = Tensor([4,5,6], dtype=dtypes.int)
out = a + b
out.callify()
self.assertListEqual(out.tolist(), [5, 7, 9])
def test_reduce(self):
out = Tensor([1.,2,3,4]).sum()
out.callify()
self.assertEqual(out.item(), 10.0)
def test_multiple_ops(self):
a = Tensor([1.,2,3])
b = Tensor([4.,5,6])
out = (a + b) * (a - b)
out.callify()
self.assertListEqual(out.tolist(), [-15.0, -21.0, -27.0])
def test_double_callify(self):
a = Tensor([1.,2,3])
b = Tensor([4.,5,6])
out = a + b
out.callify()
out.callify()
self.assertListEqual(out.tolist(), [5.0, 7.0, 9.0])
def test_double_callify_multi_output(self):
a = Tensor([1.,2,3])
b = Tensor([4.,5,6])
c = a + b
d = a * b
c.callify(d)
c.callify(d)
self.assertListEqual(c.tolist(), [5.0, 7.0, 9.0])
self.assertListEqual(d.tolist(), [4.0, 10.0, 18.0])
if __name__ == "__main__":
unittest.main()

125
artifacts/package_sources/tinygrad/test/unit/test_conv.py Normal file
View File

@@ -0,0 +1,125 @@
import unittest
import numpy as np
from tinygrad.tensor import Tensor
from tinygrad.helpers import Context
class TestConv(unittest.TestCase):
def test_simple(self):
x = Tensor.ones(1,12,16,32).contiguous().realize()
w = Tensor.ones(32,12,3,3).contiguous().realize()
ret = x.conv2d(w, stride=(2,2), padding=(1,1)).numpy()
# it's not 108 around the padding
assert (ret[:, :, 1:-1, 1:-1] == 108).all()
assert ret[0,0,0,0] == 48
assert ret[0,0,0,1] == 72
def test_simple_rand(self):
x = Tensor.rand(1,12,16,32)
w = Tensor.rand(32,12,3,3)
x.conv2d(w, stride=(2,2), padding=(1,1)).numpy()
def test_many_simple(self):
x = Tensor(np.arange(8*2*8).reshape(1,8,2,8).astype(np.float32))
#w = Tensor(np.arange(8*8*1*1).reshape(8,8,1,1).astype(np.float32))
w = Tensor.eye(8).reshape((8,8,1,1))
ret = x.conv2d(w, stride=(1,2), padding=(0,0)).numpy()
print(ret)
def test_lazycache(self):
x = Tensor.rand(1, 32)
y = Tensor.rand(32)
out = x + y.reshape((1,32,1)).reshape((1,32)) + y.reshape((1,32,1)).reshape((1,32))
out.numpy()
def test_simple_biased(self):
C = 8
x = Tensor.rand(1,C,5,5)
w = Tensor.eye(C).reshape((C,C,1,1))
b = Tensor(np.arange(C).astype(np.float32))
ret = Tensor.conv2d(x,w,b).relu().conv2d(w,b)
print(ret.numpy())
def test_two_binops_no_rerun_small(self):
x = Tensor.rand(1,1,32,32)
w = Tensor.rand(1,1,3,3)
out = x.conv2d(w, padding=(1,1))
np.testing.assert_allclose(out.relu().numpy(), np.maximum(out.numpy(), 0), atol=1e-6)
def test_two_overlapping_binops_no_rerun(self):
x = Tensor.randn(1,12,16,32)
w = Tensor.randn(32,12,3,3)
out = x.conv2d(w, stride=(2,2), padding=(1,1))
r1, r2 = out.relu(), out.elu()
np.testing.assert_allclose(r1.numpy(), np.maximum(out.numpy(), 0), atol=1e-5)
np.testing.assert_allclose(r2.numpy(), np.where(out.numpy() > 0, out.numpy(), (np.exp(out.numpy()) - 1)), atol=1e-5)
@unittest.skip("this test is flaky")
def test_two_overlapping_binops_no_rerun_wino(self):
with Context(WINO=1):
x = Tensor.randn(1,4,16,16)
w = Tensor.randn(6,4,3,3)
out = x.conv2d(w, padding=(1,1))
r1, r2 = out.relu(), out.elu()
np.testing.assert_allclose(r1.numpy(), np.maximum(out.numpy(), 0), atol=1e-5)
np.testing.assert_allclose(r2.numpy(), np.where(out.numpy() > 0, out.numpy(), (np.exp(out.numpy()) - 1)), atol=1e-5)
def test_first_three(self):
x = Tensor.rand(1,12,16,32)
w = Tensor.rand(32,12,3,3)
x = x.conv2d(w, stride=(2,2), padding=(1,1)).elu()
w = Tensor.rand(32,1,3,3)
x = x.conv2d(w, padding=(1,1), groups=32).elu()
w = Tensor.rand(16,32,1,1)
x = x.conv2d(w).elu()
x = x.numpy()
print(x.shape)
def test_elu(self):
x = Tensor.rand(1,12,16,32)
w = Tensor.rand(32,12,3,3)
x = x.conv2d(w, stride=(2,2), padding=(1,1))
x = x.elu()
w = Tensor.rand(32,1,3,3)
x = x.conv2d(w, padding=(1,1), groups=32)
x.numpy()
def test_reduce_relu(self):
x = Tensor.rand(1,12,16,32)
x = x.sum(keepdim=True).relu()
x.numpy()
def test_bias(self):
from tinygrad.nn import Conv2d
x = Tensor.rand(1,12,16,32)
c = Conv2d(12, 32, 3)
x = c(x).relu()
w = Tensor.uniform(32, 1, 3, 3)
x = x.conv2d(w, groups=32)
x.numpy()
def test_multiadd(self):
w = Tensor.rand(32)
x = Tensor.rand(32).relu()
(w+x).numpy()
def test_reorder(self):
x = Tensor.rand(1,12,16,32)
w = Tensor.rand(12,12,3,3)
x = x.conv2d(w, padding=(1,1))
print(x.shape)
x = x.reshape((1, 12, 32, 16))
x += 1
x += 1
x = x.reshape((1, 12, 16, 32))
x.numpy()
if __name__ == '__main__':
unittest.main()

26
artifacts/package_sources/tinygrad/test/unit/test_cpu.py Normal file
View File

@@ -0,0 +1,26 @@
import unittest, io
from contextlib import redirect_stdout
from tinygrad import Tensor, Device
from tinygrad.helpers import Target
from tinygrad.renderer.nir import LVPRenderer
from tinygrad.renderer.isa.x86 import X86Renderer
from tinygrad.codegen import to_program
@unittest.skipIf(Device.DEFAULT != "CPU", "only run on CPU")
class TestCPU(unittest.TestCase):
def test_arch_feats(self):
ast = (Tensor.empty(16) + Tensor.empty(16)).schedule_linear().src[-1].src[0]
for ren in Device[Device.DEFAULT].renderers:
for arch, expect_vmov in [("x86_64,x86-64,avx", True), ("x86_64,x86-64,-avx", False)]:
with self.subTest(arch=arch):
if ren is X86Renderer: continue # X86 requires avx support
if ren is LVPRenderer: continue # LVP does not play nice with cross compilation
r = ren(Target(device="CPU", arch=arch))
p = to_program(ast, r)
lib = r.compiler.compile(p.src[2].arg)
out = io.StringIO()
with redirect_stdout(out): r.compiler.disassemble(lib)
self.assertEqual("vmov" in out.getvalue(), expect_vmov, out.getvalue())
if __name__ == '__main__':
unittest.main()

579
artifacts/package_sources/tinygrad/test/unit/test_disk_tensor.py Normal file
View File

@@ -0,0 +1,579 @@
import os, pathlib, tempfile, unittest
import numpy as np
from tinygrad import Tensor, Device, dtypes
from tinygrad.dtype import DType, DTYPES_DICT
from tinygrad.nn.state import safe_load, safe_save, get_state_dict, torch_load
from tinygrad.helpers import Timing, fetch, OSX, dedup, Context
from test.helpers import slow
class TempDirTestCase(unittest.TestCase):
def setUp(self):
self.temp_dir = tempfile.TemporaryDirectory()
def tearDown(self):
self.temp_dir.cleanup()
def tmp(self, name:str) -> str:
return (pathlib.Path(self.temp_dir.name) / name).as_posix()
def compare_weights_both(url):
import torch
fn = fetch(url)
tg_weights = get_state_dict(torch_load(fn))
torch_weights = get_state_dict(torch.load(fn, map_location=torch.device('cpu'), weights_only=False), tensor_type=torch.Tensor)
assert list(tg_weights.keys()) == list(torch_weights.keys())
for k in tg_weights:
if tg_weights[k].dtype == dtypes.bfloat16: tg_weights[k] = torch_weights[k].float() # numpy doesn't support bfloat16
if torch_weights[k].dtype == torch.bfloat16: torch_weights[k] = torch_weights[k].float() # numpy doesn't support bfloat16
if torch_weights[k].requires_grad: torch_weights[k] = torch_weights[k].detach()
np.testing.assert_equal(tg_weights[k].numpy(), torch_weights[k].numpy(), err_msg=f"mismatch at {k}, {tg_weights[k].shape}")
print(f"compared {len(tg_weights)} weights")
class TestTorchLoad(TempDirTestCase):
# pytorch pkl format
def test_load_enet(self): compare_weights_both("https://github.com/lukemelas/EfficientNet-PyTorch/releases/download/1.0/efficientnet-b0-355c32eb.pth")
# pytorch zip format
def test_load_enet_alt(self): compare_weights_both("https://download.pytorch.org/models/efficientnet_b0_rwightman-3dd342df.pth")
# pytorch zip format
def test_load_convnext(self): compare_weights_both('https://dl.fbaipublicfiles.com/convnext/convnext_tiny_1k_224_ema.pth')
def test_load_llama2bfloat(self): compare_weights_both("https://huggingface.co/qazalin/bf16-lightweight/resolve/main/consolidated.00.pth?download=true")
# pytorch tar format
def test_load_resnet(self): compare_weights_both('https://download.pytorch.org/models/resnet50-19c8e357.pth')
# shared storage (mixtral-8x7b-32kseqlen)
def test_shared_storage(self):
import torch
fn = self.tmp("shared_storage.pth")
torch.save({"a": (a := torch.randn(100)), "b": a[5:]}, fn)
compare_weights_both(fn)
test_fn = pathlib.Path(__file__).parents[2] / "weights/LLaMA/7B/consolidated.00.pth"
#test_size = test_fn.stat().st_size
test_size = 1024*1024*1024*2
def _test_bitcasted(t: Tensor, dt: DType, expected):
np.testing.assert_allclose(t.bitcast(dt).numpy(), expected)
# sudo su -c 'sync; echo 1 > /proc/sys/vm/drop_caches' && python3 test/unit/test_disk_tensor.py TestRawDiskBuffer.test_readinto_read_speed
class TestRawDiskBuffer(unittest.TestCase):
@unittest.skipIf(not test_fn.exists(), "download LLaMA weights for read in speed tests")
def test_readinto_read_speed(self):
tst = np.empty(test_size, np.uint8)
with open(test_fn, "rb") as f:
with Timing("copy in ", lambda et_ns: f" {test_size/et_ns:.2f} GB/s"):
f.readinto(tst)
def test_bitcasts_on_disk(self):
_, tmp = tempfile.mkstemp()
# ground truth = https://evanw.github.io/float-toy/
t = Tensor.empty((128, 128), dtype=dtypes.uint8, device=f"disk:{tmp}") # uint8
# all zeroes
_test_bitcasted(t, dtypes.float16, 0.0)
_test_bitcasted(t, dtypes.uint16, 0)
_test_bitcasted(t, dtypes.float32, 0.0)
_test_bitcasted(t, dtypes.uint32, 0)
# pi in float16 stored via int16
t.bitcast(dtypes.uint16).assign(Tensor.full((128, 64), 0x4248, dtype=dtypes.uint16)).realize()
_test_bitcasted(t, dtypes.float16, 3.140625)
_test_bitcasted(t, dtypes.float32, 50.064727)
_test_bitcasted(t, dtypes.uint16, 0x4248)
_test_bitcasted(t, dtypes.uint32, 0x42484248)
# pi in float32 stored via float32
t.bitcast(dtypes.float32).assign(Tensor.full((128, 32), 3.1415927, dtype=dtypes.float32)).realize()
_test_bitcasted(t, dtypes.float32, 3.1415927)
_test_bitcasted(t, dtypes.uint32, 0x40490FDB)
# doesn't suport normal cast
with self.assertRaises(NotImplementedError):
Tensor.empty((4,), dtype=dtypes.int16, device=f"disk:{tmp}").cast(dtypes.float16).to(None).realize()
# Those two should be moved to test_dtype.py:test_shape_change_bitcast after bitcast works on non-disk
with self.assertRaises(RuntimeError):
# should fail because 3 int8 is 3 bytes but float16 is two and 3 isn't a multiple of 2
Tensor.empty((3,), dtype=dtypes.int8, device=f"DISK:{tmp}").bitcast(dtypes.float16).shape
pathlib.Path(tmp).unlink()
class TestSafetensors(TempDirTestCase):
def test_real_safetensors(self):
import torch
from safetensors.torch import save_file
torch.manual_seed(1337)
tensors = {
"weight1": torch.randn((16, 16)),
"weight2": torch.arange(0, 17, dtype=torch.uint8),
"weight3": torch.arange(0, 17, dtype=torch.int32).reshape(17,1,1),
"weight4": torch.arange(0, 2, dtype=torch.uint8),
}
save_file(tensors, self.tmp("real.safetensors"))
ret = safe_load(self.tmp("real.safetensors"))
for k,v in tensors.items(): np.testing.assert_array_equal(ret[k].numpy(), v.numpy())
safe_save(ret, self.tmp("real.safetensors_alt"))
with open(self.tmp("real.safetensors"), "rb") as f:
with open(self.tmp("real.safetensors_alt"), "rb") as g:
assert f.read() == g.read()
ret2 = safe_load(self.tmp("real.safetensors_alt"))
for k,v in tensors.items(): np.testing.assert_array_equal(ret2[k].numpy(), v.numpy())
def test_real_safetensors_open(self):
fn = self.tmp("real_safe")
state_dict = {"tmp": Tensor.rand(10,10)}
safe_save(state_dict, fn)
import os
assert os.path.getsize(fn) == 8+0x40+(10*10*4)
from safetensors import safe_open
with safe_open(fn, framework="pt", device="cpu") as f:
assert sorted(f.keys()) == sorted(state_dict.keys())
for k in f.keys():
np.testing.assert_array_equal(f.get_tensor(k).numpy(), state_dict[k].numpy())
@unittest.skip("this test takes 7 seconds. TODO: make disk assign lazy")
def test_efficientnet_safetensors(self):
from extra.models.efficientnet import EfficientNet
model = EfficientNet(0)
state_dict = get_state_dict(model)
safe_save(state_dict, self.tmp("eff0"))
state_dict_loaded = safe_load(self.tmp("eff0"))
assert sorted(state_dict_loaded.keys()) == sorted(state_dict.keys())
for k,v in state_dict.items():
np.testing.assert_array_equal(v.numpy(), state_dict_loaded[k].numpy())
# load with the real safetensors
from safetensors import safe_open
with safe_open(self.tmp("eff0"), framework="pt", device="cpu") as f:
assert sorted(f.keys()) == sorted(state_dict.keys())
for k in f.keys():
np.testing.assert_array_equal(f.get_tensor(k).numpy(), state_dict[k].numpy())
def _test_huggingface_enet_safetensors(self, fn):
state_dict = safe_load(fn)
assert len(state_dict.keys()) == 244
assert 'blocks.2.2.se.conv_reduce.weight' in state_dict
assert state_dict['blocks.0.0.bn1.num_batches_tracked'].numpy() == 276570
assert state_dict['blocks.2.0.bn2.num_batches_tracked'].numpy() == 276570
def test_huggingface_enet_safetensors(self):
# test a real file
fn = fetch("https://huggingface.co/timm/mobilenetv3_small_075.lamb_in1k/resolve/main/model.safetensors")
self._test_huggingface_enet_safetensors(fn)
def test_huggingface_enet_safetensors_fromurl(self):
# test tensor input
t = Tensor.from_url("https://huggingface.co/timm/mobilenetv3_small_075.lamb_in1k/resolve/main/model.safetensors")
self._test_huggingface_enet_safetensors(t)
def test_metadata(self):
metadata = {"hello": "world"}
safe_save({}, self.tmp('metadata.safetensors'), metadata)
import struct
with open(self.tmp('metadata.safetensors'), 'rb') as f:
dat = f.read()
sz = struct.unpack(">Q", dat[0:8])[0]
import json
assert json.loads(dat[8:8+sz])['__metadata__']['hello'] == 'world'
def test_safe_save_only_copy(self):
from tinygrad.helpers import GlobalCounters
t = Tensor.rand(10, 10).realize()
GlobalCounters.reset()
safe_save({"t": t}, self.tmp("test_copy.safetensors"))
assert GlobalCounters.global_ops == 0, f"safe_save should have no compute, got {GlobalCounters.global_ops} ops"
def test_save_all_dtypes(self):
for dtype in dedup(DTYPES_DICT.values()):
if dtype in dtypes.fp8_fnuz: continue # not supported by safetensors
path = self.tmp(f"ones.{dtype}.safetensors")
ones = Tensor(np.random.rand(10,10), dtype=dtype)
safe_save(get_state_dict(ones), path)
loaded = list(safe_load(path).values())[0]
# numpy has no fp8 or bfloat16, compare the stored bytes
if dtype == dtypes.bfloat16 or dtype in dtypes.fp8s:
np.testing.assert_equal(ones.bitcast(dtypes.uint8).numpy(), loaded.bitcast(dtypes.uint8).numpy())
else: np.testing.assert_equal(ones.numpy(), loaded.numpy())
def test_load_supported_types(self):
import torch
from safetensors.torch import save_file
from safetensors.numpy import save_file as np_save_file
torch.manual_seed(1337)
tensors = {
"weight_F16": torch.randn((2, 2), dtype=torch.float16),
"weight_F32": torch.randn((2, 2), dtype=torch.float32),
"weight_U8": torch.tensor([1, 2, 3], dtype=torch.uint8),
"weight_I8": torch.tensor([-1, 2, 3], dtype=torch.int8),
"weight_I32": torch.tensor([-1, 2, 3], dtype=torch.int32),
"weight_I64": torch.tensor([-1, 2, 3], dtype=torch.int64),
"weight_F64": torch.randn((2, 2), dtype=torch.double),
"weight_BOOL": torch.tensor([True, False], dtype=torch.bool),
"weight_I16": torch.tensor([127, 64], dtype=torch.short),
"weight_BF16": torch.randn((2, 2), dtype=torch.bfloat16),
}
save_file(tensors, self.tmp("dtypes_torch.safetensors"))
loaded = safe_load(self.tmp("dtypes_torch.safetensors"))
for k,v in loaded.items():
if v.dtype != dtypes.bfloat16:
assert v.numpy().dtype == tensors[k].numpy().dtype
np.testing.assert_allclose(v.numpy(), tensors[k].numpy())
# pytorch does not support U16, U32, and U64 dtypes.
tensors = {
"weight_U16": np.array([1, 2, 3], dtype=np.uint16),
"weight_U32": np.array([1, 2, 3], dtype=np.uint32),
"weight_U64": np.array([1, 2, 3], dtype=np.uint64),
}
np_save_file(tensors, self.tmp("dtypes_numpy.safetensors"))
loaded = safe_load(self.tmp("dtypes_numpy.safetensors"))
for k,v in loaded.items():
assert v.numpy().dtype == tensors[k].dtype
np.testing.assert_allclose(v.numpy(), tensors[k])
def helper_test_disk_tensor(tmp, fn, data, np_fxn, tinygrad_fxn=None):
if tinygrad_fxn is None: tinygrad_fxn = np_fxn
pathlib.Path(tmp(fn)).unlink(missing_ok=True)
tinygrad_tensor = Tensor(data, device="CPU").to(f"disk:{tmp(fn)}")
numpy_arr = np.array(data)
tinygrad_fxn(tinygrad_tensor)
np_fxn(numpy_arr)
np.testing.assert_allclose(tinygrad_tensor.numpy(), numpy_arr)
class TestDiskTensor(TempDirTestCase):
def test_empty(self):
Tensor.empty(100, 100, device=f"disk:{self.tmp('dt_empty')}")
def test_simple_read(self):
fn = pathlib.Path(self.tmp("dt_simple_read"))
fn.write_bytes(bytes(range(256)))
t = Tensor.empty(16, 16, device=f"disk:{self.tmp('dt_simple_read')}", dtype=dtypes.uint8)
out = t[1].to(Device.DEFAULT).tolist()
assert out == list(range(16, 32))
def test_simple_read_bitcast(self):
fn = pathlib.Path(self.tmp("dt_simple_read_bitcast"))
fn.write_bytes(bytes(range(256))*2)
t = Tensor.empty(16, 16*2, device=f"disk:{self.tmp('dt_simple_read_bitcast')}", dtype=dtypes.uint8)
out = t[1].bitcast(dtypes.uint16).to(Device.DEFAULT).tolist()
tout = [(x//256, x%256) for x in out]
assert tout == list([(x+1,x) for x in range(32,64,2)])
def test_simple_read_bitcast_alt(self):
fn = pathlib.Path(self.tmp("dt_simple_read_bitcast_alt"))
fn.write_bytes(bytes(range(256))*2)
t = Tensor.empty(16, 16*2, device=f"disk:{self.tmp('dt_simple_read_bitcast_alt')}", dtype=dtypes.uint8)
out = t.bitcast(dtypes.uint16)[1].to(Device.DEFAULT).tolist()
tout = [(x//256, x%256) for x in out]
assert tout == list([(x+1,x) for x in range(32,64,2)])
def test_strided_read(self):
# test non-contiguous (strided) read raises
dt = Tensor([0, 1, 2, 3, 4, 5]).to(f"disk:{self.tmp('dt_strided_read')}")
with self.assertRaisesRegex(RuntimeError, "non-contiguous view is not supported"):
dt[::2].tolist()
def test_permuted_read(self):
# test non-contiguous (permuted) read raises
dt = Tensor([[0, 1, 2], [3, 4, 5]]).to(f"disk:{self.tmp('dt_permuted_read')}")
with self.assertRaisesRegex(RuntimeError, "non-contiguous view is not supported"):
dt.T.tolist()
def test_write_ones(self):
out = Tensor.ones(10, 10, device="CPU").contiguous()
outdisk = out.to(f"disk:{self.tmp('dt_write_ones')}")
print(outdisk)
outdisk.realize()
del out, outdisk
import struct
# test file
with open(self.tmp("dt_write_ones"), "rb") as f:
assert f.read() == struct.pack('<f', 1.0) * 100 == b"\x00\x00\x80\x3F" * 100
# test load alt
reloaded = Tensor.empty(10, 10, device=f"disk:{self.tmp('dt_write_ones')}")
np.testing.assert_almost_equal(reloaded.numpy(), np.ones((10, 10)))
def test_simple_setitem(self):
data = [[1],[2]]
src = Tensor(data)
dt = src.to(f"disk:{self.tmp('dt_simple_setitem')}")
dt[1] = [3]
self.assertEqual(dt.tolist(), [[1], [3]])
def test_strided_setitem(self):
# test non-contiguous (strided) setitem raises
dt = Tensor([1, 2, 3, 4, 5, 6]).to(f"disk:{self.tmp('dt_strided_setitem')}")
with self.assertRaisesRegex(RuntimeError, "non-contiguous view is not supported"):
dt[::2] = Tensor([10, 20, 30])
def test_advanced_setitem_not_supported(self):
dt = Tensor.arange(12).reshape(3, 4).clone().to(f"disk:{self.tmp('dt_advanced_setitem')}")
with self.assertRaises(RuntimeError, msg="advanced setitem is not supported for DISK tensors"):
dt[Tensor([0, 2]), Tensor([1, 3])] = 99
def test_assign_const_to_disk(self):
# assign from CONST (Tensor.full) to disk - source has no buffer, needs contiguous first
dt = Tensor.empty(4, device=f"disk:{self.tmp('dt_assign_const')}", dtype=dtypes.int32)
dt.assign(Tensor.full((4,), 42, dtype=dtypes.int32)).realize()
np.testing.assert_array_equal(dt.numpy(), [42, 42, 42, 42])
def test_assign_slice_from_const(self):
# slice assign from CONST to disk - tests size calculation when no RANGE ops
dt = Tensor([0, 1, 2, 3], dtype=dtypes.int32).to(f"disk:{self.tmp('dt_slice_const')}")
dt[1:3].assign(Tensor.full((2,), 99, dtype=dtypes.int32)).realize()
np.testing.assert_array_equal(dt.numpy(), [0, 99, 99, 3])
def test_disk_to_disk_copy(self):
# disk-to-disk copy needs to go through CPU
src = Tensor([1, 2, 3, 4], dtype=dtypes.int32).to(f"disk:{self.tmp('dt_d2d_src')}")
dst = Tensor.empty(4, device=f"disk:{self.tmp('dt_d2d_dst')}", dtype=dtypes.int32)
dst.assign(src.to("CPU")).realize()
np.testing.assert_array_equal(dst.numpy(), [1, 2, 3, 4])
def test_assign_slice(self):
def assign(x,s,y): x[s] = y
helper_test_disk_tensor(self.tmp, "dt_assign_slice_1", [0,1,2,3], lambda x: assign(x, slice(0,2), [13, 12]))
helper_test_disk_tensor(self.tmp, "dt_assign_slice_2", [[0,1,2,3],[4,5,6,7]], lambda x: assign(x, slice(0,1), [[13, 12, 11, 10]]))
def test_reshape(self):
helper_test_disk_tensor(self.tmp, "dt_reshape_1", [1,2,3,4,5], lambda x: x.reshape((1,5)))
helper_test_disk_tensor(self.tmp, "dt_reshape_2", [1,2,3,4], lambda x: x.reshape((2,2)))
def test_assign_to_different_dtype(self):
# NOTE: this is similar to Y_train in fetch_cifar
t = Tensor.empty(10, device=f'disk:{self.tmp("dt_assign_to_different_dtype")}', dtype=dtypes.int64)
for i in range(5):
data = np.array([3, 3])
idx = 2 * i
t[idx:idx+2].assign(data)
np.testing.assert_array_equal(t.numpy(), np.array([3] * 10))
def test_assign_with_bitcast(self):
# bitcast assign is used in safe_save for writing header length
t = Tensor.empty(16, device=f"disk:{self.tmp('dt_assign_bitcast')}", dtype=dtypes.uint8)
t[0:8].bitcast(dtypes.int64).assign([12345])
val = int.from_bytes(t[0:8].data(), 'little')
self.assertEqual(val, 12345)
def test_assign_to_bitcast_view(self):
# assign float values to a float32 view of a uint8 disk buffer (used by safe_save)
t = Tensor.empty(32, device=f"disk:{self.tmp('dt_bitcast_view_assign')}", dtype=dtypes.uint8)
# create float32 view of bytes 8-24 (4 floats)
float_view = t[8:24].bitcast(dtypes.float32)
float_view.assign(Tensor([1.0, 2.0, 3.0, 4.0], dtype=dtypes.float32, device="CPU"))
np.testing.assert_array_equal(float_view.numpy(), [1.0, 2.0, 3.0, 4.0])
def test_assign_cross_device(self):
# disk assign allows cross-device (source on GPU/CPU, target on disk)
t = Tensor.empty(4, device=f"disk:{self.tmp('dt_assign_cross')}", dtype=dtypes.float32)
src = Tensor([1.0, 2.0, 3.0, 4.0]) # on default device
t.assign(src)
np.testing.assert_array_equal(t.numpy(), [1.0, 2.0, 3.0, 4.0])
def test_bitcast(self):
with open(self.tmp('dt_bitcast'), "wb") as f: f.write(bytes(range(10,20)))
t = Tensor.empty(5, dtype=dtypes.int16, device=f"disk:{self.tmp('dt_bitcast')}")
ret = t.to("CPU").bitcast(dtypes.uint16) + 1
assert ret.tolist() == [2827, 3341, 3855, 4369, 4883]
def test_bitcast_view(self):
with open(self.tmp('dt_bitcast_view'), "wb") as f: f.write(bytes(range(10, 24)))
t = Tensor.empty(3, dtype=dtypes.uint, device=f"disk:{self.tmp('dt_bitcast_view')}").shrink([(0, 2)])
ret = t.bitcast(dtypes.uint16).to("CPU") + 1
assert ret.tolist() == [2827, 3341, 3855, 4369]
@unittest.skipIf(OSX or Device.DEFAULT == "CL", "new LLVM has an issue on OSX, DEV=CL gives the wrong output")
def test_bf16_disk_write_read(self):
t = Tensor([10000, -1, -1000, -10000, 20], dtype=dtypes.float32)
t.to(f"disk:{self.tmp('dt_bf16_disk_write_read_f32')}").realize()
# hack to "cast" f32 -> bf16
with open(self.tmp('dt_bf16_disk_write_read_f32'), "rb") as f: dat = f.read()
adat = b''.join([dat[i+2:i+4] for i in range(0, len(dat), 4)])
with open(self.tmp('dt_bf16_disk_write_read_bf16'), "wb") as f: f.write(adat)
t = Tensor.empty(5, dtype=dtypes.bfloat16, device=f"disk:{self.tmp('dt_bf16_disk_write_read_bf16')}")
ct = t.to(Device.DEFAULT).cast(dtypes.float)
assert ct.numpy().tolist() == [9984., -1, -1000, -9984, 20]
def test_copy_from_disk(self):
fn = pathlib.Path(self.tmp("dt_copy_from_disk"))
fn.write_bytes(bytes(range(256))*1024)
t = Tensor.empty(256*1024, device=f"disk:{self.tmp('dt_copy_from_disk')}", dtype=dtypes.uint8)
on_dev = t.to(Device.DEFAULT).realize()
np.testing.assert_equal(on_dev.numpy(), t.numpy())
def test_copy_from_disk_offset(self):
fn = pathlib.Path(self.tmp("dt_copy_from_disk_offset"))
fn.write_bytes(bytes(range(256))*1024)
for off in [314, 991, 2048, 4096]:
t = Tensor.empty(256*1024, device=f"disk:{self.tmp('dt_copy_from_disk_offset')}", dtype=dtypes.uint8)[off:]
on_dev = t.to(Device.DEFAULT).realize()
np.testing.assert_equal(on_dev.numpy(), t.numpy())
def test_shard_copy_from_disk_slice(self):
fn = pathlib.Path(self.tmp("dt_shard_copy_from_disk_slice"))
fn.write_bytes(bytes(range(32)))
with Context(CACHELEVEL=0):
t = Tensor.empty(8, 4, device=f"disk:{fn}", dtype=dtypes.uint8)[0:4].shard(("CPU:0", "CPU:1"), axis=0).realize()
np.testing.assert_equal(t.to("CPU").numpy(), np.arange(16, dtype=np.uint8).reshape(4, 4))
@slow
def test_copy_from_disk_huge(self):
fn = pathlib.Path(self.tmp("dt_copy_from_disk_huge"))
fn.write_bytes(bytes(range(256))*1024*256)
for off in [0, 551]:
t = Tensor.empty(256*1024*256, device=f"disk:{self.tmp('dt_copy_from_disk_huge')}", dtype=dtypes.uint8)[off:]
on_dev = t.to(Device.DEFAULT).realize()
np.testing.assert_equal(on_dev.numpy(), t.numpy())
@unittest.skip("this allocates a lot of RAM")
@unittest.skipUnless(OSX, "seems to only be an issue on macOS with file size >2 GiB")
def test_copy_to_cpu_not_truncated(self):
fn = self.tmp("dt_copy_to_cpu_not_truncated")
with open(fn, "wb") as f: f.write(b'\x01' * (size := int(2 * 1024**3)) + (test := b"test"))
x = Tensor.empty(size + len(test), dtype=dtypes.uint8, device=f"disk:{fn}").to("CPU").realize()
assert x[size:].data().tobytes() == test
def test_disk_device_reuse(self):
from tinygrad.runtime.ops_disk import DiskDevice
fn = pathlib.Path(self.tmp("dt_device_reuse"))
fn.write_bytes(bytes(range(256)))
# create first tensor and realize it
t1 = Tensor.empty(128, device=f"disk:{fn}", dtype=dtypes.uint8)
t1.to("CPU").realize()
# get the DiskDevice and check internal state
disk_device = Device[f"DISK:{fn}"]
assert isinstance(disk_device, DiskDevice)
assert disk_device.refcount == 1
assert hasattr(disk_device, "mem")
first_fd = disk_device.fd
# create second tensor on same file - should reuse the device, not re-open
t2 = Tensor.empty(64, device=f"disk:{fn}", dtype=dtypes.uint8)
t2.to("CPU").realize()
assert disk_device.refcount == 2
assert disk_device.fd == first_fd, "file descriptor changed - file was unnecessarily re-opened"
# verify data is correct
np.testing.assert_equal(t1.numpy(), np.arange(128, dtype=np.uint8))
np.testing.assert_equal(t2.numpy(), np.arange(64, dtype=np.uint8))
@unittest.skip("fails with setup_python_cap run")
def test_disk_open_failure_state(self):
from tinygrad.runtime.ops_disk import DiskDevice
fn = pathlib.Path(self.tmp("dt_open_failure"))
fn.write_bytes(bytes(range(256)))
os.chmod(fn, 0o000)
try:
t = Tensor.empty(100, device=f"disk:{fn}", dtype=dtypes.uint8)
t.numpy()
except PermissionError: pass
# device state should be clean after failed open
disk_device = Device[f"DISK:{fn}"]
assert isinstance(disk_device, DiskDevice)
assert disk_device.size is None, "size should be None after failed open"
assert not hasattr(disk_device, "mem"), "mem should not exist after failed open"
# should be able to open with any size after failure
os.chmod(fn, 0o644)
t2 = Tensor.empty(200, device=f"disk:{fn}", dtype=dtypes.uint8)
t2.to("CPU").realize()
assert disk_device.size == 200
@unittest.skip("fails with setup_python_cap run")
def test_disk_permission_error(self):
fn = pathlib.Path(self.tmp("dt_permission"))
fn.write_bytes(bytes(range(256)))
os.chmod(fn, 0o000)
try:
with self.assertRaises(PermissionError):
Tensor.empty(100, device=f"disk:{fn}", dtype=dtypes.uint8).numpy()
finally:
os.chmod(fn, 0o644)
class TestPathTensor(TempDirTestCase):
def setUp(self):
super().setUp()
self.test_file = pathlib.Path(self.temp_dir.name) / "test_file.bin"
self.test_data = np.arange(100, dtype=np.uint8).tobytes()
with open(self.test_file, "wb") as f:
f.write(self.test_data)
def test_path_tensor_no_device(self):
t = Tensor(self.test_file)
self.assertEqual(t.shape, (100,))
self.assertEqual(t.dtype, dtypes.uint8)
self.assertTrue(t.device.startswith("DISK:"))
np.testing.assert_array_equal(t.numpy(), np.frombuffer(self.test_data, dtype=np.uint8))
def test_path_tensor_with_device(self):
t = Tensor(self.test_file, device="CPU")
self.assertEqual(t.shape, (100,))
self.assertEqual(t.dtype, dtypes.uint8)
self.assertEqual(t.device, "CPU")
np.testing.assert_array_equal(t.numpy(), np.frombuffer(self.test_data, dtype=np.uint8))
def test_path_tensor_empty_file(self):
empty_file = pathlib.Path(self.temp_dir.name) / "empty_file.bin"
empty_file.touch()
t = Tensor(empty_file)
self.assertEqual(t.shape, (0,))
self.assertEqual(t.dtype, dtypes.uint8)
self.assertTrue(t.device.startswith("DISK:"))
def test_path_tensor_non_existent_file(self):
non_existent_file = pathlib.Path(self.temp_dir.name) / "non_existent.bin"
with self.assertRaises(FileNotFoundError):
Tensor(non_existent_file)
def test_path_tensor_with_dtype(self):
t = Tensor(self.test_file, dtype=dtypes.int16)
self.assertEqual(t.shape, (50,))
self.assertEqual(t.dtype, dtypes.int16)
self.assertTrue(t.device.startswith("DISK:"))
np.testing.assert_array_equal(t.numpy(), np.frombuffer(self.test_data, dtype=np.int16))
def test_path_tensor_copy_to_device(self):
t = Tensor(self.test_file)
t_cpu = t.to("CPU")
self.assertEqual(t_cpu.device, "CPU")
np.testing.assert_array_equal(t_cpu.numpy(), np.frombuffer(self.test_data, dtype=np.uint8))
@unittest.skip("permission checks don't work in all environments")
def test_path_tensor_disk_device_bug(self):
test_file = pathlib.Path(self.temp_dir.name) / "disk_device_bug"
with open(test_file, "wb") as f: f.write(bytes(range(10)))
os.chmod(test_file, 0o000)
with self.assertRaises(PermissionError):
Tensor(pathlib.Path(test_file)).tolist()
os.chmod(test_file, 0o644)
assert Tensor(pathlib.Path(test_file)).tolist(), list(range(10))
class TestDiskTensorMovement(TempDirTestCase):
def setUp(self):
super().setUp()
self.fn = pathlib.Path(self.tmp("custom_disk_range"))
Tensor.arange(100, dtype=dtypes.uint8).clone().to(f"disk:{str(self.fn)}").realize()
def test_simple_read(self):
t = Tensor(self.fn)
self.assertTrue(Tensor.all(t.to(None) == Tensor.arange(100, dtype=dtypes.uint8)).item())
def test_slice_read(self):
t = Tensor(self.fn)
self.assertListEqual(t[16:18].tolist(), [16,17])
def test_slice_read_cat(self):
t = Tensor(self.fn)
with self.assertRaises(AssertionError):
self.assertListEqual(Tensor.cat(t[16:18], t[20:22]).tolist(), [16,17,20,21])
def test_slice_sum(self):
t = Tensor(self.fn)
with self.assertRaises(AssertionError):
self.assertListEqual((t[16:18]+t[20:22]).tolist(), [16+20,17+21])
if __name__ == "__main__":
unittest.main()

243
artifacts/package_sources/tinygrad/test/unit/test_dtype_spec.py Normal file
View File

@@ -0,0 +1,243 @@
import unittest, math, subprocess
from tinygrad.tensor import Tensor
from tinygrad.dtype import dtypes, DType, DTYPES_DICT, strong_dtype
from tinygrad.device import Device
from tinygrad.helpers import getenv, DEBUG, EMULATED_DTYPES, Context
from test.helpers import slow
from hypothesis import given, settings, strategies as strat
import numpy as np
import torch
settings.register_profile("my_profile", max_examples=50, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
settings.load_profile("my_profile")
supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes()
core_dtypes = list(DTYPES_DICT.values())
dtype_ints = [dt for dt in core_dtypes if dtypes.is_int(dt) and dt in supported_dtypes]
dtype_floats = [dt for dt in core_dtypes if dtypes.is_float(dt) and dt in supported_dtypes]
FP8E4M3_MAX = 448.0
FP8E5M2_MAX = 57344.0
FP8E4M3FNUZ_MAX = 240.0
FP8E5M2FNUZ_MAX = 57344.0
def _assert_eq(tensor:Tensor, target_dtype:DType, target, tol_target_dtype:float=1e-7):
if DEBUG >= 2: print(tensor.numpy())
try:
assert tensor.dtype == target_dtype
# weak values read back at their default.
target_dtype = strong_dtype(target_dtype)
# denormals are zero
if target_dtype in dtypes.floats and (target_dtype not in supported_dtypes or target_dtype in EMULATED_DTYPES.tolist(dtypes)):
fe, fm = dtypes.finfo(target_dtype)
kwargs = {"atol":2 ** (2 - (1 << (fe - 1))), "rtol": 2 ** (-fm)}
else: kwargs = {"rtol": {dtypes.float16:1e-3, dtypes.bfloat16:1e-2, dtypes.fp8e4m3:1e-1, dtypes.fp8e5m2:5e-1,
dtypes.fp8e4m3fnuz:1e-1, dtypes.fp8e5m2fnuz:5e-1}.get(target_dtype, tol_target_dtype)}
np.testing.assert_allclose(tensor.numpy(), target, **kwargs)
except AssertionError as e:
raise AssertionError(f"\ntensor {tensor.numpy()} dtype {tensor.dtype} does not match target {target} with dtype {target_dtype}") from e
class TestTypeSpec(unittest.TestCase):
def test_default_dtype_context(self):
default_float, default_int = dtypes.default_float, dtypes.default_int
with Context(DEFAULT_FLOAT=dtypes.half, DEFAULT_INT=dtypes.int16):
assert dtypes.default_float is dtypes.half
assert dtypes.default_int is dtypes.int16
assert dtypes.default_float is default_float
assert dtypes.default_int is default_int
@unittest.skip("this test is slow and spawning whole pythons")
def test_env_set_default_float(self):
# check default
subprocess.run(['python3 -c "from tinygrad import dtypes; assert dtypes.default_float == dtypes.float"'],
shell=True, check=True)
# check change
subprocess.run(['DEFAULT_FLOAT=HALF python3 -c "from tinygrad import dtypes; assert dtypes.default_float == dtypes.half"'],
shell=True, check=True)
# check invalid
with self.assertRaises(subprocess.CalledProcessError):
subprocess.run(['DEFAULT_FLOAT=INT32 python3 -c "from tinygrad import dtypes"'],
shell=True, check=True)
with self.assertRaises(subprocess.CalledProcessError):
subprocess.run(['DEFAULT_FLOAT=TYPO python3 -c "from tinygrad import dtypes"'],
shell=True, check=True)
def test_dtype_str_arg(self):
n = np.random.normal(0, 1, (10, 10)).astype(np.float32)
tested = 0
for dtype_str, dtype in [
("bool", dtypes.bool), ("int8", dtypes.int8), ("int", dtypes.int), ("uint32", dtypes.uint32), ("float32", dtypes.float32)]:
with np.errstate(invalid='ignore'):
np.testing.assert_equal(Tensor(n, dtype=dtype_str).numpy(), Tensor(n, dtype=dtype).numpy())
np.testing.assert_equal(Tensor(n).cast(dtype_str).numpy(), Tensor(n).cast(dtype).numpy())
if dtype.itemsize == 4:
np.testing.assert_equal(Tensor(n).bitcast(dtype_str).numpy(), Tensor(n).bitcast(dtype).numpy())
tested += 1
assert tested == 3
with self.assertRaises(AttributeError): Tensor([1, 2, 3], dtype="nonexistdtype")
with self.assertRaises(AttributeError): Tensor([1, 2, 3], dtype="")
np.testing.assert_equal(Tensor(n).sum(dtype="int16").numpy(), Tensor(n).sum(dtype=dtypes.int16).numpy())
@given(strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
def test_creation(self, default_int, default_float):
self.enterContext(Context(DEFAULT_INT=default_int, DEFAULT_FLOAT=default_float))
_assert_eq(Tensor(True), dtypes.bool, True)
_assert_eq(Tensor(None), dtypes.weakfloat, [])
_assert_eq(Tensor(2), dtypes.weakint, 2)
_assert_eq(Tensor(2.34), dtypes.weakfloat, 2.34)
_assert_eq(Tensor([]), dtypes.default_float, [])
_assert_eq(Tensor([1]), dtypes.default_int, [1])
# list elements are python scalars; a numpy scalar in a list has no inferred dtype (use np.array or state a dtype)
with self.assertRaises(RuntimeError): Tensor([np.int32(1)])
_assert_eq(Tensor([1.1]), dtypes.default_float, [1.1])
_assert_eq(Tensor.eye(0), dtypes.default_float, np.eye(0))
_assert_eq(Tensor.eye(3), dtypes.default_float, np.eye(3))
_assert_eq(Tensor.eye(3, dtype=dtypes.int64), dtypes.int64, np.eye(3))
if dtypes.float16 in supported_dtypes:
_assert_eq(Tensor.eye(3, dtype=dtypes.float16), dtypes.float16, np.eye(3))
@given(strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
def test_full(self, default_int, default_float):
self.enterContext(Context(DEFAULT_INT=default_int, DEFAULT_FLOAT=default_float))
_assert_eq(Tensor.zeros((2, 3)), dtypes.default_float, np.zeros((2, 3)))
_assert_eq(Tensor.zeros((2, 3), dtype=dtypes.int64), dtypes.int64, np.zeros((2, 3)))
if dtypes.float16 in supported_dtypes:
_assert_eq(Tensor.zeros((2, 3), dtype=dtypes.float16), dtypes.float16, np.zeros((2, 3)))
_assert_eq(Tensor.ones((2, 3)), dtypes.default_float, np.ones((2, 3)))
_assert_eq(Tensor.ones((2, 3), dtype=dtypes.int64), dtypes.int64, np.ones((2, 3)))
if dtypes.float16 in supported_dtypes:
_assert_eq(Tensor.ones((2, 3), dtype=dtypes.float16), dtypes.float16, np.ones((2, 3)))
_assert_eq(Tensor.full((2, 3), 3.0), dtypes.default_float, np.full((2, 3), 3.0))
_assert_eq(Tensor.full((2, 3), 3), dtypes.default_int, np.full((2, 3), 3))
_assert_eq(Tensor.full((2, 3), True), dtypes.bool, np.full((2, 3), True))
_assert_eq(Tensor.full((2, 3), 3, dtype=dtypes.int64), dtypes.int64, np.full((2, 3), 3))
_assert_eq(Tensor.full((2, 3), 3.0, dtype=dtypes.int64), dtypes.int64, np.full((2, 3), 3))
if dtypes.float16 in supported_dtypes:
_assert_eq(Tensor.full((2, 3), 3, dtype=dtypes.float16), dtypes.float16, np.full((2, 3), 3))
_assert_eq(Tensor.full((2, 3), 3.0, dtype=dtypes.float16), dtypes.float16, np.full((2, 3), 3))
@given(strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
def test_reduce_0d_default(self, default_int, default_float):
self.enterContext(Context(DEFAULT_INT=default_int, DEFAULT_FLOAT=default_float))
_assert_eq(Tensor.ones((2,3,0)).sum(2), dtypes.default_float, np.zeros((2, 3)))
# TODO: what should this one be?
# _assert_eq(Tensor.ones((2,3,0), dtype=dtypes.default_int).sum(2), dtypes.default_int, np.zeros((2, 3)))
_assert_eq(Tensor.ones((2,3,0), dtype=dtypes.int32).sum(2), dtypes.int32, np.zeros((2, 3)))
@given(strat.sampled_from(dtype_ints), strat.sampled_from(dtype_floats))
def test_arange(self, default_int, default_float):
self.enterContext(Context(DEFAULT_INT=default_int, DEFAULT_FLOAT=default_float))
_assert_eq(Tensor.arange(5), dtypes.default_int, np.arange(5))
_assert_eq(Tensor.arange(120), dtypes.default_int, np.arange(120))
_assert_eq(Tensor.arange(5.0), dtypes.default_float, np.arange(5))
if dtypes.int16 in supported_dtypes:
_assert_eq(Tensor.arange(5, dtype=dtypes.int16), dtypes.int16, np.arange(5))
_assert_eq(Tensor.arange(5, dtype=dtypes.int64), dtypes.int64, np.arange(5))
if dtypes.float16 in supported_dtypes:
_assert_eq(Tensor.arange(5, dtype=dtypes.float16), dtypes.float16, np.arange(5))
_assert_eq(Tensor.arange(3, 9, 0.7), dtypes.default_float, np.arange(3, 9, 0.7), 1e-6 if Device.DEFAULT == "WEBGPU" else 1e-7)
_assert_eq(Tensor.arange(3, 8.5, 3), dtypes.default_float, np.arange(3, 8.5, 3))
# stop-start and step have different signs
_assert_eq(Tensor.arange(3, 5, -2), dtypes.default_int, np.arange(3, 5, -2))
_assert_eq(Tensor.arange(5.0, 3.0), dtypes.default_float, np.arange(5.0, 3.0))
class TestAutoCastType(unittest.TestCase):
def test_int_sqrt(self):
_assert_eq(Tensor([1, 4, 9, 16]).sqrt(), dtypes.default_float, [1, 2, 3, 4])
@given(strat.sampled_from([d for d in core_dtypes if dtypes.is_int(d) and d in supported_dtypes]))
def test_int_to_float_unary_func(self, dtype):
for func in [
lambda t: t.exp(),
lambda t: t.exp2(),
lambda t: t.log(),
lambda t: t.log2(),
lambda t: t.sqrt(),
lambda t: t.rsqrt(),
lambda t: t.sin(),
lambda t: t.cos(),
lambda t: t.tan(),
lambda t: t.sigmoid(),
]:
a = [2, 3, 4]
# float16 can have larger precision errors
np.testing.assert_allclose(func(Tensor(a, dtype=dtype)).numpy(), func(torch.tensor(a)), rtol=1e-3, atol=1e-3)
@unittest.skipUnless(dtypes.float16 in supported_dtypes, "need float16")
def test_sum_dtype_arg(self):
t = Tensor([40000, 40000], dtype=dtypes.float16)
# default float16 sum returns in float16, overflowed in this case
assert t.sum().dtype == dtypes.float16
assert math.isinf(t.sum().numpy().item())
# specifiying dtype and it's not downcasted
assert t.sum(dtype=dtypes.float32).dtype == dtypes.float32
np.testing.assert_allclose(t.sum(dtype=dtypes.float32).numpy(), 80000)
def test_prod_dtype_arg(self):
t = Tensor([100, 200], dtype=dtypes.int32)
assert t.prod().dtype == dtypes.int32
np.testing.assert_allclose(t.prod().numpy(), 20000)
assert t.prod(dtype=dtypes.float32).dtype == dtypes.float32
np.testing.assert_allclose(t.prod(dtype=dtypes.float32).numpy(), 20000)
def test_gradient_dtype(self):
for default_dtype in dtypes.floats:
if default_dtype not in supported_dtypes: continue
with Context(DEFAULT_FLOAT=default_dtype):
for dtype in dtypes.floats:
if dtype not in supported_dtypes: continue
if DEBUG >= 2:
print(f"testing {default_dtype=}, {dtype=}")
a = Tensor([1, 2, 3], dtype=dtype)
b = (a * 5).sum()
b.backward() # if there is dtype mismatch, lazy should assert
assert a.grad.dtype == a.dtype
np.testing.assert_allclose(a.grad.numpy(), [5, 5, 5])
@unittest.skipIf(Device.DEFAULT == "PYTHON", "very slow")
@slow
@unittest.skipIf(Device.DEFAULT == "WEBGPU", "Binding size is larger than the maximum storage buffer binding size")
@unittest.skipUnless(dtypes.half in supported_dtypes, "need half")
def test_mean_half_precision_underflow(self):
N = 10000
x = 0.001
t = Tensor([[x]], dtype=dtypes.half).expand(N, N).contiguous()
np.testing.assert_allclose(t.mean(axis=1).numpy(), np.array([x] * N, dtype=np.float16), rtol=1e-3)
@unittest.skip("this test only works with SPLIT_REDUCEOP=1")
@unittest.skipUnless(dtypes.half in supported_dtypes, "need half")
def test_mean_half_precision_overflow(self):
N = 256
t = Tensor([60000] * N*N, dtype=dtypes.half).reshape(N, N)
np.testing.assert_allclose(t.mean().numpy(), 60000)
t.square().mean().backward()
np.testing.assert_allclose(t.grad.numpy().flatten(), [60000 * 2 / (N*N)] * N*N)
@unittest.skipIf(Device.DEFAULT == "WEBGPU", "Precision error")
@unittest.skipUnless(dtypes.half in supported_dtypes, "need half")
def test_softmax_dtype(self):
data = [1, 2, 3]
t = Tensor(data, dtype=dtypes.half)
tt = torch.tensor(data, dtype=torch.half)
out = t.softmax(0)
self.assertEqual(out.dtype, dtypes.half)
np.testing.assert_allclose(out.numpy(), tt.softmax(0).numpy(), rtol=1e-3)
out = t.softmax(0, dtype=dtypes.float)
self.assertEqual(out.dtype, dtypes.float)
np.testing.assert_allclose(out.numpy(), tt.softmax(0, dtype=torch.float).numpy(), rtol=1e-3)
out = t.log_softmax(0)
self.assertEqual(out.dtype, dtypes.half)
np.testing.assert_allclose(out.numpy(), tt.log_softmax(0).numpy(), rtol=1e-3)
out = t.log_softmax(0, dtype=dtypes.float)
self.assertEqual(out.dtype, dtypes.float)
np.testing.assert_allclose(out.numpy(), tt.log_softmax(0, dtype=torch.float).numpy(), rtol=1e-3)

280
artifacts/package_sources/tinygrad/test/unit/test_dtype_weak.py Normal file
View File

@@ -0,0 +1,280 @@
import tempfile, unittest, math
from tinygrad import Tensor, dtypes, TinyJit
from tinygrad.helpers import Context
from tinygrad.dtype import least_upper_float
from tinygrad.uop.ops import UOp, Ops, dtype_from_uop, graph_rewrite
from tinygrad.uop.weak import pm_lower_index_dtype, pm_commit_weak
from tinygrad.uop.symbolic import symbolic_simple
from tinygrad.uop.spec import spec_shared, type_verify
from tinygrad.engine.jit import JitError
class TestWeakPromotion(unittest.TestCase):
def test_rand_requires_concrete(self):
with self.assertRaises(ValueError): Tensor.rand(2, dtype=dtypes.weakfloat)
with self.assertRaises(ValueError): Tensor.const(1.0).rand_like()
with self.assertRaises(ValueError): Tensor.const(1.0).randn_like()
def test_reduce_strips_weakness(self):
for weak, value, strong in ((dtypes.weakint, 1, dtypes.default_int), (dtypes.weakfloat, 1.0, dtypes.default_float)):
t = Tensor.const(value, weak).expand(3)
for out in (t.sum(), t.max(), t.prod(), t.cumsum(0), t.cummax(0)[0]): self.assertEqual(out.dtype, strong)
self.assertEqual((Tensor.const(1.0).expand(3).sum() + Tensor([1], dtype=dtypes.float16)).dtype, dtypes.float32)
def test_materialize_at_default_dtype(self):
for weak, value, strong in ((dtypes.weakfloat, 0.5, dtypes.default_float),):
t = Tensor.const(value, weak)
self.assertEqual(t.dtype, weak)
self.assertEqual(t.data().itemsize, strong.itemsize)
self.assertEqual(t.numpy().dtype.itemsize, strong.itemsize)
# materializing commits at the kind default; contiguous has no layout to fix so it stays weak
self.assertEqual((c := t.clone("CPU")).dtype, strong)
self.assertEqual(c.item(), value)
self.assertEqual(t.contiguous().dtype, weak)
def test_assign_into_weak_commits(self):
t = Tensor.const(0.5)
t.assign(Tensor(1.0, dtype=dtypes.default_float))
self.assertEqual((t.dtype, t.item()), (dtypes.default_float, 1.0))
def test_float_unary_on_weakint_stays_weak(self):
self.assertIs(least_upper_float(dtypes.weakint), dtypes.weakfloat)
def test_copysign_meets_operands(self):
r = Tensor([2], dtype=dtypes.uint8, device="CPU").copysign(Tensor([1], dtype=dtypes.uint32, device="CPU"))
self.assertEqual((r.dtype, r.tolist()), (dtypes.uint32, [2]))
def test_minimum_reflects_weak_operand(self):
r = Tensor(1).minimum(Tensor([2], dtype=dtypes.uint8, device="CPU"))
self.assertEqual((r.dtype, r.tolist()), (dtypes.uint8, [1]))
for dt in dtypes.uints:
r = Tensor([dt.max], dtype=dt, device="CPU").minimum(1)
self.assertEqual((r.dtype, r.tolist()), (dt, [1]))
self.assertNotIn(Ops.CAST, [u.op for u in r._uop.toposort()])
def test_broadcasted_keeps_const_weak(self):
# a python scalar stays a bare weak CONST through _broadcasted, lifted only to the KIND of the lub
x, y = Tensor([1], dtype=dtypes.int8)._broadcasted(3)
self.assertEqual((y._uop.base.op, y.dtype, x.dtype), (Ops.CONST, dtypes.weakint, dtypes.int8))
x, y = Tensor([1], dtype=dtypes.int8)._broadcasted(0.5)
self.assertEqual((y._uop.base.op, y.dtype, x.dtype), (Ops.CONST, dtypes.weakfloat, dtypes.weakfloat))
x, y = Tensor.const(1).reshape(1)._broadcasted(Tensor([1.0], dtype=dtypes.float32))
self.assertEqual((x._uop.base.op, x._uop.base.val, x.dtype, x.shape, y.dtype),
(Ops.CONST, 1, dtypes.weakfloat, (1,), dtypes.float32))
def test_weak_expression_anchors_at_strong_lub(self):
# regression test for the HALF bert nan (#17408, reverted in #17409): lub(int32, weakfloat)==weakfloat makes
# `loss_mask.sum() + 1e-5` a weakfloat EXPRESSION. Meeting a strong float in a binop must pin it at the lub
denom = (Tensor.zeros(912, dtype=dtypes.int32) != Tensor.zeros(912, dtype=dtypes.float32)).sum() + 1e-5
self.assertIs(denom.dtype, dtypes.weakfloat) # the setup: the denominator expression itself is weak
x, y = Tensor([2048.0], dtype=dtypes.float32)._broadcasted(denom)
self.assertIs(y.dtype, dtypes.float32)
recips = [u for u in (x / y)._uop.toposort() if u.op is Ops.RECIPROCAL]
self.assertEqual([(u.dtype, u.src[0].dtype) for u in recips], [(dtypes.float32, dtypes.float32)])
with Context(DEFAULT_FLOAT=dtypes.float16):
committed = graph_rewrite((UOp.const(1).cast(dtypes.int32) + UOp.const(1.0)).cast(dtypes.float32), pm_lower_index_dtype, ctx={})
self.assertEqual([u.dtype for u in committed.toposort() if u.op is Ops.ADD], [dtypes.float32])
def test_cast_weak_expression_commits_at_cast_floor(self):
# the floor never narrows: a cast BELOW the default does not pull the compute width down with it
with Context(DEFAULT_FLOAT=dtypes.float32):
narrowed = graph_rewrite((UOp.const(1.0) + UOp.const(2.0)).cast(dtypes.float16), pm_lower_index_dtype, ctx={})
self.assertEqual((narrowed.dtype, narrowed.src[0].dtype), (dtypes.float16, dtypes.float32))
def test_cast_weak_expression_value_uses_cast_floor(self):
with Context(DEFAULT_FLOAT=dtypes.float16):
denom = Tensor.ones(1, dtype=dtypes.int32, device="CPU").sum() * 70000 + 1e-5
out = Tensor(1.0, dtype=dtypes.float32, device="CPU") / denom
self.assertAlmostEqual(out.item(), 1 / (70000 + 1e-5), places=10)
def test_uop_scalar_const_lifts_kind(self):
for dtype, value, out_dtype, const_dtype in ((dtypes.weakint, 1, dtypes.weakint, dtypes.weakint),
(dtypes.int32, 1, dtypes.int32, dtypes.weakint),
(dtypes.int32, 0.5, dtypes.weakfloat, dtypes.weakfloat),
(dtypes.float32, 1, dtypes.float32, dtypes.weakfloat)):
out = UOp.variable("x", 0.0 if dtype == dtypes.float32 else 0, 10.0 if dtype == dtypes.float32 else 10, dtype) + value
self.assertEqual((out.dtype, out.src[1].op, out.src[1].dtype), (out_dtype, Ops.CONST, const_dtype))
# the kind lift converts the VALUE too (the arg is the only dtype carrier once UOp.const loses its dtype arg),
# and a bare weak const UOp is the same spelling as the python scalar: both lift to the same node
x = UOp.variable("x", 0.0, 1.0, dtypes.float32)
self.assertIsInstance((x + 2).src[1].val, float)
self.assertIs(x + UOp.const(2), x + 2)
def test_index_dtype_ignores_weakness(self):
with Context(SPEC=2):
idx = UOp.const(0).cast(dtypes.int32)
weak = UOp.const(1.0).expand((1,))
self.assertEqual(UOp(Ops.INDEX, dtypes.float32, (weak, idx)).dtype, dtypes.float32)
with self.assertRaisesRegex(RuntimeError, "bad dtype"): UOp(Ops.INDEX, dtypes.int32, (weak, idx))
def test_store_weak_value_uses_destination_dtype(self):
with Context(DEFAULT_FLOAT=dtypes.float16):
dst = UOp.param(0, dtypes.bfloat16, (1,)).index(UOp.const(0).cast(dtypes.int32))
gate = UOp.const(True)
out = graph_rewrite(dst.store(UOp.const(5.0), gate), pm_lower_index_dtype, ctx={})
# a bare weak CONST commits directly: the pass runs without symbolic, so a CAST here would survive it
self.assertEqual((out.src[1], out.src[2]), (UOp.const(5.0, dtypes.bfloat16), gate))
def test_weak_srcs_commit_only_at_a_concrete_lub(self):
weak_lub = UOp(Ops.ADD, src=(UOp.const(1), UOp.const(1.0)))
self.assertIs(graph_rewrite(weak_lub, pm_lower_index_dtype, ctx={}), weak_lub)
concrete = UOp.const(2.0).cast(dtypes.float16)
where = graph_rewrite(UOp(Ops.WHERE, src=(UOp.const(True), concrete, UOp.const(1.0))), pm_lower_index_dtype, ctx={})
self.assertEqual(tuple(x.dtype for x in where.src), (dtypes.bool, dtypes.float16, dtypes.float16))
def test_weak_shift_lhs_commits_the_node(self):
# a shift derives its lhs's dtype, so committing the lhs restates the root (WGSL's packed store writes `mask << shift_am`)
shl = graph_rewrite(UOp.const(0xFFFF) << UOp.variable("x", 0, 16, dtypes.uint), symbolic_simple+pm_commit_weak)
self.assertEqual((shl.dtype, shl.src[0]), (dtypes.uint, UOp.const(0xFFFF, dtypes.uint)))
@unittest.expectedFailure # TODO: a weak const defers to its consumer (JAX): these dtypes change once python scalars are weak consts
def test_changed_rows(self):
t_i8, t_f16, t_bf16 = Tensor([1], dtype=dtypes.int8), Tensor([1], dtype=dtypes.float16), Tensor([1], dtype=dtypes.bfloat16)
t_bool, t_u16 = Tensor([True]), Tensor([1], dtype=dtypes.uint16)
self.assertEqual((t_i8 + 0.5).dtype, dtypes.weakfloat)
self.assertEqual(((t_i8 + 0.5) + t_f16).dtype, dtypes.float16)
self.assertEqual(((t_i8 + 0.5) + t_bf16).dtype, dtypes.bfloat16)
self.assertEqual(((t_bool + 1) + t_i8).dtype, dtypes.int8)
self.assertEqual(((t_bool + 1) + t_u16).dtype, dtypes.uint16)
self.assertEqual((Tensor(3) + t_i8).dtype, dtypes.int8)
# zeros/ones are full with a python fill value, so they are weak too (jnp.zeros pins float32; deliberate divergence)
self.assertEqual((Tensor.zeros(3) + t_f16).dtype, dtypes.float16)
def test_unchanged_rows(self):
t_i8, t_f16, t_f32 = Tensor([1], dtype=dtypes.int8), Tensor([1], dtype=dtypes.float16), Tensor([1], dtype=dtypes.float32)
self.assertEqual((t_i8 + 1).dtype, dtypes.int8)
self.assertEqual((t_f16 + 0.5).dtype, dtypes.float16)
self.assertEqual((t_f32 + t_f16).dtype, dtypes.float32)
self.assertEqual(Tensor([2], dtype=dtypes.uint8).pad(((1, 1),), value=1).dtype, dtypes.uint8)
def test_concrete_pair_promotes_weak(self):
out = Tensor([-1], dtype=dtypes.int64, device="CPU") + Tensor([3], dtype=dtypes.uint64, device="CPU") + Tensor(0.5)
self.assertEqual((out.dtype, out.tolist()), (dtypes.weakfloat, [2.5]))
def test_dot_defers_weak(self):
weak = Tensor([True, False]).where(Tensor(1), 2)
self.assertEqual(weak.dot(Tensor([1, 1], dtype=dtypes.int8)).dtype, dtypes.int8)
def test_weak_int_binop(self):
v = UOp.variable("i", 0, 10, dtypes.weakint)
self.assertEqual((v << 1).dtype, dtypes.weakint)
self.assertEqual(dtype_from_uop(Ops.SHL, (UOp.const(1, dtypes.int8), UOp.const(1, dtypes.uint32)), None), dtypes.int8)
self.assertEqual(UOp.const(1).alu(Ops.SHL, UOp.const(1, dtypes.uint)).dtype, dtypes.weakint)
self.assertEqual((v & 3).dtype, dtypes.weakint)
with self.assertRaises(RuntimeError): Tensor.const(1.0) << Tensor.const(1.0)
with self.assertRaises(RuntimeError): UOp.const(1, dtypes.int32).alu(Ops.SHL, UOp.const(1, dtypes.float64))
for op in (Ops.SHL, Ops.SHR):
with self.assertRaises(RuntimeError):
UOp.const(1, dtypes.float32).alu(op, UOp.const(1, dtypes.int32))
# float bitwise builds, the spec rejects it
with Context(SPEC=1):
f32, wf = UOp.const(1.0, dtypes.float32), UOp.const(1.0)
for bad in (f32.alu(Ops.AND, f32), UOp(Ops.AND, dtypes.float32, (f32, f32)), UOp(Ops.AND, dtypes.int32, (wf, wf))):
with self.assertRaises(RuntimeError): type_verify([bad], spec_shared)
def test_integer_values(self):
x = Tensor.full((1,), 1, dtype=dtypes.int64, device="CPU")
self.assertEqual((x + 2**40).item(), 2**40 + 1)
self.assertEqual((x << 3).item(), 8)
self.assertTrue((x < 2**40).item())
def test_float64_precision(self):
value = 1.0 + 2**-40
x64 = Tensor.full((1,), 1.0, dtype=dtypes.float64, device="CPU")
self.assertEqual((x64 + value).item(), 2.0 + 2**-40)
x32 = Tensor.full((1,), 0.0, dtype=dtypes.float32, device="CPU")
self.assertEqual((x32 + value).item(), 1.0)
def test_weak_transcendentals(self):
t_f16 = Tensor([1], dtype=dtypes.float16)
for out in (Tensor(2).exp(), Tensor(2).cos(), Tensor(2).sigmoid()):
self.assertEqual((out.dtype, (out + t_f16).dtype), (dtypes.weakfloat, dtypes.float16))
def test_null_lowering(self):
for t in (Tensor.full((1,), 1, dtype=dtypes.int64, device="NULL") + 2**40,
Tensor.full((1,), 1.0, dtype=dtypes.float64, device="NULL") + (1.0 + 2**-40)):
t.realize()
self.assertNotIn(t.uop.buffer.dtype, dtypes.weaks)
def test_computed_float_index_lowers(self):
# a half-pixel nearest index resolves its float-scaled range before the gather
idx = (Tensor.arange(8) + 0.5) / 4 - 0.5
idx = (idx.clip(0, 1) - 0.5).ceil().int()
out = Tensor([0, 1], device="NULL")[idx].contiguous().realize()
self.assertNotIn(out.uop.buffer.dtype, dtypes.weaks)
class TestWeakStorageBoundary(unittest.TestCase):
# weak has no storage: a weak assignment source casts when it defers to the destination, everything else raises
def test_weak_source(self):
w05 = Tensor.const(0.5).reshape(1)
dst = Tensor.zeros(2, dtype=dtypes.int8, device="CPU").contiguous().realize()
with self.assertRaises(RuntimeError): dst.assign(w05.expand(2)) # weakfloat into int does not defer
with self.assertRaises(RuntimeError): dst[0:1] = w05
fdst = Tensor.zeros(2, dtype=dtypes.float32, device="CPU").contiguous().realize()
fdst[0:1] = w05 # weakfloat defers to float
self.assertEqual(fdst.tolist(), [0.5, 0.0])
with tempfile.TemporaryDirectory() as td: # the DISK path checks the same
ddst = Tensor.empty(2, dtype=dtypes.int32, device=f"DISK:{td}/t")
with self.assertRaises(RuntimeError): ddst.assign(w05.expand(2))
def test_weak_has_no_storage(self):
import numpy as np
with self.assertRaises(RuntimeError): Tensor(np.ones(2, dtype=np.float32), dtype=dtypes.weakfloat)
with self.assertRaises(RuntimeError): Tensor(bytes(8), dtype=dtypes.weakfloat)
class TestWeakMaterializationEntries(unittest.TestCase):
# everything that creates storage from a weak value raises
def test_reads_commit_storage_raises(self):
for weak, value, strong in ((dtypes.weakfloat, 0.5, dtypes.default_float),):
def weak_val():
return Tensor([True], device="CPU").where(Tensor.const(value, weak), Tensor.const(value, weak))
self.assertEqual(weak_val().dtype, weak)
self.assertEqual(weak_val().to("CPU").dtype, weak)
self.assertEqual(weak_val().data().format, strong.fmt)
self.assertEqual(weak_val().numpy().dtype.itemsize, strong.itemsize)
self.assertEqual(weak_val().tolist(), [value])
self.assertEqual(weak_val().cast(strong).realize().uop.buffer.dtype, strong)
self.assertEqual(weak_val().contiguous().dtype, weak) # no layout to fix, stays weak
self.assertEqual(weak_val().realize().dtype, weak) # no width to store, stays weak
self.assertEqual(weak_val().clone().dtype, strong) # storage commits at the default
for entry in (lambda t: t.to("CPU:1").realize(), lambda t: t.as_param(0)):
with self.assertRaises(RuntimeError): entry(weak_val())
def test_weak_is_virtual(self):
# NOTE: int64 lub uint64 is weakfloat, so this is device-ful weak from promotion, never from a cast to weak
devful = Tensor([1], dtype=dtypes.int64, device="CPU") + Tensor([1], dtype=dtypes.uint64, device="CPU")
for t in (Tensor.const(0.5), devful):
self.assertTrue(t.uop.is_virtual)
# realize is a no-op, so a weak input can never become the real buffer TinyJit needs
with self.assertRaises(JitError): TinyJit(lambda x: (x+1).realize())(t)
# callify must not silently commit a weak CONTIGUOUS to storage
c = devful.alu(Ops.CONTIGUOUS)
c.callify()
self.assertIs(c.dtype, dtypes.weakfloat)
def test_empty_reads_commit(self):
for weak, strong in ((dtypes.weakfloat, dtypes.default_float),):
empty = Tensor.const(0, weak).reshape(1).shrink(((0, 0),))
self.assertEqual(empty.data().format, strong.fmt)
self.assertEqual(empty.numpy().dtype.itemsize, strong.itemsize)
self.assertEqual(empty.tolist(), [])
class TestSignedUint64Weakfloat(unittest.TestCase):
# int64 and uint64 have no common integer supertype (JAX JEP), so the join defers to weakfloat instead of wrapping
def test_no_wrap(self):
r = Tensor([-1], dtype=dtypes.int64, device="CPU") + Tensor([1], dtype=dtypes.uint64, device="CPU")
self.assertEqual((r.dtype, r.item()), (dtypes.weakfloat, 0.0))
def test_weakfloat_lowers(self):
i64, u64 = Tensor([-1], dtype=dtypes.int64, device="CPU"), Tensor([3], dtype=dtypes.uint64, device="CPU")
r = i64 + u64 + Tensor([2], dtype=dtypes.float16, device="CPU") # a concrete consumer takes the join
self.assertEqual((r.dtype, r.cast(dtypes.float32).item()), (dtypes.half, 4.0))
self.assertEqual((i64 < u64).item(), True) # comparison meets at float
self.assertAlmostEqual((i64 + u64).sin().item(), math.sin(2), places=5) # Unary lowers before transcendental
if __name__ == "__main__":
unittest.main()

654
artifacts/package_sources/tinygrad/test/unit/test_function.py Normal file
View File

@@ -0,0 +1,654 @@
import numpy as np
import unittest
from tinygrad.function import function
from tinygrad import Tensor, GlobalCounters, Device
from tinygrad.dtype import Invalid
from tinygrad.uop.ops import UOp, Ops, KernelInfo, ProgramInfo
from test.helpers import assert_kernel_count
class TestFunction(unittest.TestCase):
def test_simple(self):
@function
def f(a:Tensor, b:Tensor) -> Tensor: return a+b
a = Tensor([1,2,3])
b = Tensor([4,5,6])
np.testing.assert_equal(f(a,b).numpy(), [5,7,9])
def test_simple_same(self):
@function
def f(a:Tensor, b:Tensor) -> Tensor: return a+b
a = Tensor([1,2,3])
np.testing.assert_equal(f(a,a).numpy(), [2,4,6])
def test_depth_restored_on_exception(self):
from tinygrad.function import _function
@function
def f(a:Tensor) -> Tensor: raise ValueError("error")
with self.assertRaises(ValueError): f(Tensor([1]))
self.assertEqual(_function.depth, 0)
def test_implicit(self):
inp = Tensor([7,8,9])
@function(allow_implicit=True)
def f(a:Tensor, b:Tensor) -> Tensor: return a+b+inp
a = Tensor([1,2,3])
b = Tensor([4,5,6])
np.testing.assert_equal(f(a,b).numpy(), [12,15,18])
def test_implicit_same_as_input(self):
inp = Tensor([7,8,9])
@function(allow_implicit=True)
def f(a:Tensor, b:Tensor) -> Tensor: return a+b+inp
a = Tensor([1,2,3])
np.testing.assert_equal(f(a, inp).numpy(), [15,18,21])
def test_implicit_2(self):
inp = Tensor([7,8,9])
@function(allow_implicit=True)
def f(a:Tensor, b:Tensor) -> Tensor:
return a+b+inp
inp2 = Tensor([7,8,10])
@function(allow_implicit=True)
def g(a:Tensor, b:Tensor) -> Tensor:
return a+b+inp2
a = Tensor([1,2,3])
b = Tensor([4,5,6])
c = f(a,b)
d = g(a,b)
c.realize(d)
np.testing.assert_equal(c.numpy(), [12,15,18])
np.testing.assert_equal(d.numpy(), [12,15,19])
def test_implicit_unrealized(self):
inp = Tensor([1,2,3]) + Tensor([4,5,6])
@function(allow_implicit=True)
def f(a:Tensor) -> Tensor: return a + inp
np.testing.assert_equal(f(Tensor([10,20,30])).numpy(), [15,27,39])
def test_detach(self):
@function
def f(a:Tensor, b:Tensor) -> Tensor: return a.detach() + b
a = Tensor([1,2,3])
b = Tensor([4,5,6])
np.testing.assert_equal(f(a, b).numpy(), [5,7,9])
def test_contiguous_backward(self):
@function
def f(a:Tensor, b:Tensor) -> Tensor: return (a + b).contiguous_backward()
a = Tensor([1,2,3])
b = Tensor([4,5,6])
np.testing.assert_equal(f(a, b).numpy(), [5,7,9])
def test_method(self):
class Foo:
def __init__(self): self.w = Tensor([10,20,30])
@function
def __call__(self, x:Tensor) -> Tensor: return x + self.w
foo = Foo()
np.testing.assert_equal(foo(Tensor([1,2,3])).numpy(), [11,22,33])
def test_grad_gemm(self):
@function
def f(a:Tensor, b:Tensor) -> Tensor: return a @ b
a = Tensor([[1.,2.],[3.,4.]])
b = Tensor([[5.,6.],[7.,8.]])
(f(a, b).contiguous() * b).sum().backward()
Tensor.realize(a, b, a.grad, b.grad)
# L = sum((a@b) * b), dL/d(a@b) = b, dL/da = b @ b^T, dL/db = a^T @ b + (a@b)
na, nb = a.numpy(), b.numpy()
np.testing.assert_allclose(a.grad.numpy(), nb @ nb.T)
np.testing.assert_allclose(b.grad.numpy(), na.T @ nb + na @ nb)
def test_grad_implicit(self):
w = Tensor([1., 2., 3.])
w.realize() # TODO: this is required
@function(allow_implicit=True)
def f(x:Tensor) -> Tensor: return x * w
x = Tensor([4., 5., 6.])
f(x).sum().backward()
np.testing.assert_allclose(w.grad.numpy(), [4., 5., 6.])
def test_symbolic_index(self):
table = Tensor([10,20,30,40]).contiguous().realize()
@function(allow_implicit=True)
def f(x:Tensor, start_pos:int|UOp) -> Tensor:
return x + table[start_pos]
v = UOp.variable("start_pos", 0, 3)
np.testing.assert_equal(f(Tensor([1,2,3]), v.bind(0)).numpy(), [11,12,13])
def test_symbolic_shape_input(self):
table = Tensor([10,20,30,40]).contiguous().realize()
@function
def f(x:Tensor) -> Tensor: return x * 2
sz = UOp.variable("sz", 1, 3)
slic = table[:sz.bind(2)]
np.testing.assert_equal(f(slic)[:2].numpy(), [20,40])
def test_nested_calls(self):
w = Tensor([10., 20., 30.])
@function(allow_implicit=True)
def f(a:Tensor) -> Tensor: return a + w
@function(allow_implicit=True)
def g(a:Tensor) -> Tensor: return a * w
a = Tensor([1., 2., 3.])
np.testing.assert_allclose(g(f(a)).numpy(), [110., 440., 990.])
def test_nested_calls_backward(self):
w = Tensor([[1., 2.], [3., 4.]]).contiguous().realize()
@function(allow_implicit=True)
def inner(x:Tensor) -> Tensor: return x + w
@function(allow_implicit=True)
def outer(a:Tensor, b:Tensor) -> Tensor: return inner(a.reshape(1,2) + b.reshape(1,2))
a = Tensor([1., 2.])
b = Tensor([3., 4.])
outer(a, b).sum().backward()
np.testing.assert_allclose(a.grad.numpy(), [2., 2.])
np.testing.assert_allclose(b.grad.numpy(), [2., 2.])
def test_unused_param_backward(self):
@function
def f(a:Tensor, b:Tensor, c:Tensor) -> Tensor: return a + c # b is unused
a = Tensor([1., 2., 3.])
b = Tensor([4., 5., 6.])
c = Tensor([7., 8., 9.])
f(a, b, c).sum().backward()
np.testing.assert_allclose(a.grad.numpy(), [1., 1., 1.])
np.testing.assert_allclose(b.grad.numpy(), [0., 0., 0.])
np.testing.assert_allclose(c.grad.numpy(), [1., 1., 1.])
def test_name(self):
@function
def f(a:Tensor) -> Tensor: return a + 1
assert f(Tensor([1])).uop.src[0].arg.name.endswith("f")
def test_method_name(self):
class Foo:
@function
def __call__(self, x:Tensor) -> Tensor: return x + 1
assert Foo()(Tensor([1])).uop.src[0].arg.name.endswith("Foo.__call__")
def test_callable_instance(self):
class Foo:
def __init__(self): self.w = Tensor([10,20,30])
def __call__(self, x:Tensor) -> Tensor: return x + self.w
foo = Foo()
f = function(foo, allow_implicit=True)
np.testing.assert_equal(f(Tensor([1,2,3])).numpy(), [11,22,33])
assert f(Tensor([1,2,3])).uop.src[0].arg.name.endswith("Foo")
def test_iadd(self):
@function
def f(x:Tensor) -> Tensor:
x += 1
return x
a = Tensor([1,2,3]).realize()
np.testing.assert_equal(f(a).numpy(), [2,3,4])
np.testing.assert_equal(a.numpy(), [3,4,5]) # TODO: should be [1,2,3]
def test_implicit_assign(self):
a = Tensor([1,2,3])
a += 1
c = Tensor([2,2,2]).contiguous()
@function
def f(b:Tensor) -> Tensor: return a+b+c
b = Tensor([10,20,30]).realize()
np.testing.assert_equal(f(b).numpy(), [14,25,36])
def test_assign_input(self):
@function
def f(a:Tensor, b:Tensor) -> Tensor:
a.assign(b+1)
return a
a = Tensor([1,2,3]).realize()
b = Tensor([10,20,30]).realize()
np.testing.assert_equal(f(a,b).numpy(), [11,21,31])
np.testing.assert_equal(a.numpy(), [11,21,31]) # TODO: should be [1,2,3]
np.testing.assert_equal(b.numpy(), [10,20,30])
def test_view_assign_explicit_buffer(self):
"""view assign on an explicit param's buffer should not create implicit inputs."""
class State:
def __init__(self): self.buf = Tensor.zeros(2, 4).contiguous().realize()
@function(allow_implicit=False)
def __call__(self, x:Tensor) -> Tensor:
self.buf[:, 0:2].assign(x)
return self.buf[:, 0:2]
s = State()
np.testing.assert_equal(s(Tensor([[5., 6.], [7., 8.]])).numpy(), [[5., 6.], [7., 8.]])
def test_single_after_store(self):
"""AFTER(buf, STORE(view, data)) should write data through the view into buf, same as the double-after pattern."""
@function
def f(buf:Tensor, x:Tensor, start_pos:int|UOp) -> Tensor:
slice_uop = buf[:, start_pos:start_pos+1].uop
assigned = Tensor(buf.uop.after(slice_uop.store(x.uop)))
return assigned
buf = Tensor.zeros(2, 8).contiguous().realize()
x = Tensor([[1.], [2.]]).realize()
v = UOp.variable("sp", 0, 7)
r0 = f(buf, x, v.bind(0)).numpy()
np.testing.assert_equal(r0, [[1.,0.,0.,0.,0.,0.,0.,0.], [2.,0.,0.,0.,0.,0.,0.,0.]])
def test_single_after_store_precompile(self):
"""precompiled AFTER(buf, STORE(view, data)) should return buf after the store."""
@function(precompile=True)
def f(buf:Tensor, x:Tensor, start_pos:int|UOp) -> Tensor:
slice_uop = buf[:, start_pos:start_pos+1].uop
assigned = Tensor(buf.uop.after(slice_uop.store(x.uop)))
return assigned
x = Tensor([[1.], [2.]]).realize()
v = UOp.variable("sp", 0, 7)
for sp in (0, 2):
with self.subTest(sp=sp):
buf = Tensor.zeros(2, 8).clone().realize()
expected = np.zeros((2, 8), dtype=np.float32)
expected[:, sp] = [1., 2.]
np.testing.assert_equal(f(buf, x, v.bind(sp)).numpy(), expected)
np.testing.assert_equal(buf.numpy(), expected)
@unittest.expectedFailure
def test_assign_slice(self):
@function
def f(a:Tensor, b:Tensor) -> Tensor:
a[1:] = b[1:]+1
return a
a = Tensor([1,2,3]).realize()
b = Tensor([10,20,30]).realize()
np.testing.assert_equal(f(a,b).numpy(), [1,21,31])
np.testing.assert_equal(a.numpy(), [1,2,3])
np.testing.assert_equal(b.numpy(), [10,20,30])
class TestFunctionMulti(unittest.TestCase):
devices_2 = ("CPU:0", "CPU:1")
def test_simple_multi(self):
@function
def f(a:Tensor, b:Tensor) -> Tensor: return a+b
a = Tensor([1,2,3,4]).shard(self.devices_2, axis=None)
b = Tensor([10,20,30,40]).shard(self.devices_2, axis=None)
np.testing.assert_equal(f(a,b).numpy(), [11,22,33,44])
def test_simple_multi_sharded(self):
@function
def f(a:Tensor, b:Tensor) -> Tensor: return a+b
a = Tensor([1,2,3,4]).shard(self.devices_2, axis=0)
b = Tensor([10,20,30,40]).shard(self.devices_2, axis=0)
np.testing.assert_equal(f(a,b).numpy(), [11,22,33,44])
def test_data_parallel_multi(self):
@function
def f(x:Tensor, w:Tensor) -> Tensor: return x @ w
x = Tensor([[1.,2.],[3.,4.],[5.,6.],[7.,8.]]).shard(self.devices_2, axis=0)
w = Tensor([[1.,0.],[0.,1.]]).shard(self.devices_2, axis=None)
np.testing.assert_allclose(f(x, w).numpy(), [[1.,2.],[3.,4.],[5.,6.],[7.,8.]])
def test_grad_implicit_multi(self):
w = Tensor([1., 2., 3., 4.]).shard(self.devices_2, axis=None)
w.realize()
@function(allow_implicit=True)
def f(x:Tensor) -> Tensor: return x * w
x = Tensor([4., 5., 6., 7.]).shard(self.devices_2, axis=None)
f(x).sum().backward()
np.testing.assert_allclose(w.grad.numpy(), [4., 5., 6., 7.])
def test_call_axis(self):
@function
def f(x:Tensor, w:Tensor) -> Tensor: return x @ w
x = Tensor([[1.,0.],[0.,1.],[1.,1.],[0.,0.]]).shard(self.devices_2, axis=0)
w = Tensor([[1.,2.],[3.,4.]]).shard(self.devices_2, axis=None)
result = f(x, w)
# CALL output should inherit axis=0 from the sharded input
self.assertEqual(result.uop.axis, 0)
# reduce on the sharded axis should remove it
self.assertIsNone(result.sum().uop.axis)
def test_call_axis_shard_inside(self):
@function
def f(x:Tensor, w:Tensor) -> Tensor:
return x.shard(self.devices_2, axis=0) @ w.shard(self.devices_2, axis=None)
x = Tensor([[1.,0.],[0.,1.],[1.,1.],[0.,0.]])
w = Tensor([[1.,2.],[3.,4.]])
result = f(x, w)
self.assertEqual(result.uop.axis, 0)
np.testing.assert_allclose(result.numpy(), x.numpy() @ w.numpy())
def test_data_parallel_backward(self):
@function
def f(x:Tensor, w:Tensor) -> Tensor: return x @ w
x = Tensor([[1.,0.],[0.,1.],[1.,1.],[0.,0.]]).shard(self.devices_2, axis=0)
w = Tensor([[1.,2.],[3.,4.]]).shard(self.devices_2, axis=None)
w.realize()
f(x, w).sum().backward()
# d/dx = ones @ w^T = [[1,3],[1,3],[1,3],[1,3]], but sum so ones(4,2) @ w^T? no:
# L = sum(x @ w), dL/dx = ones(4,2) @ w^T... actually dL/d(xw) = ones(4,2), dL/dx = ones(4,2) @ w^T
np.testing.assert_allclose(x.grad.numpy(), np.ones((4,2)) @ np.array([[1,3],[2,4]]))
def test_data_parallel_backward_4(self):
devices_4 = tuple(f"CPU:{i}" for i in range(4))
@function
def f(x:Tensor, w:Tensor) -> Tensor: return x @ w
x = Tensor(np.arange(16).reshape(8,2).astype(np.float32)).shard(devices_4, axis=0)
w = Tensor([[1.,2.],[3.,4.]]).shard(devices_4, axis=None)
w.realize()
f(x, w).sum().backward()
np.testing.assert_allclose(x.grad.numpy(), np.ones((8,2)) @ np.array([[1,3],[2,4]]))
def test_data_parallel_backward_implicit(self):
devices_4 = tuple(f"CPU:{i}" for i in range(4))
w = Tensor([[1.,2.],[3.,4.]]).shard(devices_4, axis=None)
w.realize()
@function(allow_implicit=True)
def f(x:Tensor) -> Tensor: return x @ w
x = Tensor(np.arange(16).reshape(8,2).astype(np.float32)).shard(devices_4, axis=0)
f(x).sum().backward()
np.testing.assert_allclose(x.grad.numpy(), np.ones((8,2)) @ np.array([[1,3],[2,4]]))
def test_data_parallel_backward_twice(self):
devices_4 = tuple(f"CPU:{i}" for i in range(4))
w = Tensor([[1.,2.],[3.,4.]]).shard(devices_4, axis=None)
w.realize()
# pre-init grads like the training loop does
w.grad = w.zeros_like().contiguous().realize()
@function(allow_implicit=True)
def f(x:Tensor) -> Tensor: return x @ w
expected = np.ones((8,2)) @ np.array([[1,3],[2,4]])
for _ in range(2):
x = Tensor(np.arange(16).reshape(8,2).astype(np.float32)).shard(devices_4, axis=0)
f(x).sum().backward()
np.testing.assert_allclose(x.grad.numpy(), expected)
class TestFunctionTuple(unittest.TestCase):
def test_tuple(self, precompile=False):
x = Tensor.ones(3).contiguous()
@function(precompile=precompile)
def f(t:Tensor): return (t+1, t+2)
t1, t2 = f(x)
t1.realize(t2)
print(t1.tolist(), t2.tolist())
assert t1.tolist() == [2,2,2]
assert t2.tolist() == [3,3,3]
def test_tuple_precompile(self): self.test_tuple(True)
def test_grad_tuple(self, precompile=False):
x = Tensor.ones(3).contiguous()
y = Tensor.ones(3).contiguous()
@function(precompile=precompile)
def f(u1:Tensor, u2:Tensor): return (u1+1, u2+2)
t1, t2 = f(x,y)
(t1+t2).sum().backward()
x.grad.realize(y.grad)
def test_grad_tuple_precompile(self): self.test_grad_tuple(True)
def test_grad_fxn_tuple(self):
# grad_fxn for tuple: receives one gradient per output as positional args
def grad_fxn(d_out0:UOp, d_out1:UOp, call:UOp):
# f(u1, u2) = (u1+1, u2+2)
# df/du1 = d_out0, df/du2 = d_out1
return (d_out0, d_out1)
x = Tensor.ones(3).contiguous()
y = Tensor.ones(3).contiguous()
@function(grad_fxn=grad_fxn)
def f(u1:Tensor, u2:Tensor): return (u1+1, u2+2)
t1, t2 = f(x, y)
(t1+t2).sum().backward()
np.testing.assert_allclose(x.grad.numpy(), [1., 1., 1.])
np.testing.assert_allclose(y.grad.numpy(), [1., 1., 1.])
def test_grad_unused_tuple_output_recursive(self):
# only one output is used
@function(precompile=True, precompile_backward=True)
def f(x:Tensor, w:Tensor):
a = x @ w
b = (x @ w) * 2 # shares x@w with a
return (a, b)
x = Tensor([[1., 2.], [3., 4.]]).contiguous()
w = Tensor([[1., 0.], [0., 1.]]).contiguous()
Tensor.realize(x, w)
t1, _ = f(x, w)
t1.sum().backward()
Tensor.realize(x.grad, w.grad)
# only t1 = x @ w flows to loss; dL/dw = x.T @ ones(2,2)
np.testing.assert_allclose(w.grad.numpy(), np.array([[1., 2.], [3., 4.]]).T @ np.ones((2, 2)))
np.testing.assert_allclose(x.grad.numpy(), np.ones((2, 2)) @ np.array([[1., 0.], [0., 1.]]).T)
def test_custom_kernel_save_unused_output(self):
def my_kernel(C:UOp, D:UOp, A:UOp) -> UOp:
i = UOp.range(A.shape[0], 0)
j = UOp.range(D.shape[0], 1)
store_c = C[i].store(A[i] * 2.0).end(i)
store_d = D[j].store(A[j]).end(j)
return UOp.sink(store_c, store_d, arg=KernelInfo(name="my_kernel"))
def my_grad(d_c:UOp, call:UOp):
a_input = call.src[3]
return (None, None, (Tensor(d_c) * 2.0 + Tensor(a_input) * 0).uop)
@function(precompile=True, precompile_backward=True)
def f(a:Tensor):
c = Tensor.invalids(*a.shape, dtype=a.dtype, device=a.device)
d = Tensor.invalids(3, dtype=a.dtype, device=a.device)
c, d = Tensor.custom_kernel(c, d, a, fxn=my_kernel, grad_fxn=my_grad)[:2]
return c, d
a = Tensor([1., 2., 3., 4.]).contiguous()
Tensor.realize(a)
c, _ = f(a)
c.sum().backward()
Tensor.realize(a.grad)
np.testing.assert_allclose(a.grad.numpy(), [2., 2., 2., 2.])
def test_custom_kernel_both_outputs_used(self):
def my_kernel(C:UOp, D:UOp, A:UOp) -> UOp:
i = UOp.range(A.shape[0], 0)
store_c = C[i].store(A[i] * 2.0)
store_d = D[i].store(A[i] * 3.0)
return UOp.group(store_c, store_d).end(i).sink(arg=KernelInfo(name="my_kernel"))
def my_grad(d_c:UOp, d_d:UOp, call:UOp):
return (None, None, (Tensor(d_c) + Tensor(d_d)).uop)
@function(precompile=True, precompile_backward=True)
def f(a:Tensor):
c = Tensor.invalids(*a.shape, dtype=a.dtype, device=a.device)
d = Tensor.invalids(*a.shape, dtype=a.dtype, device=a.device)
c, d = Tensor.custom_kernel(c, d, a, fxn=my_kernel, grad_fxn=my_grad)[:2]
return (c, d)
a = Tensor([1., 2., 3., 4.]).contiguous()
Tensor.realize(a)
c, d = f(a)
(c.sum() + d.sum()).backward() # dL/da = (1 + 1) since grad_fxn passes d_combined through
Tensor.realize(a.grad)
np.testing.assert_allclose(a.grad.numpy(), [2., 2., 2., 2.])
def test_custom_kernel_precompile_no_copy_kernel(self):
def my_kernel(C:UOp, A:UOp) -> UOp:
i = UOp.range(A.shape[0], 0)
return C[i].store(A[i] * 2.0).end(i).sink(arg=KernelInfo(name="my_kernel"))
def my_grad(d_c:UOp, call:UOp):
return (None, (Tensor(d_c) * 2.0).uop)
@function(precompile=True, precompile_backward=True)
def f(a:Tensor):
c = Tensor.invalids(*a.shape, dtype=a.dtype, device=a.device)
c = Tensor.custom_kernel(c, a, fxn=my_kernel, grad_fxn=my_grad)[0]
return c
def count_kernels(t:Tensor):
linear, _ = t.linear_with_vars()
return sum((len(call.device) if isinstance(call.device, tuple) else 1)
for call in linear.src if call.src[0].op is Ops.SINK)
a = Tensor([1., 2., 3., 4.]).contiguous()
Tensor.realize(a)
c = f(a)
self.assertEqual(count_kernels(c), 1)
c.sum().backward()
Tensor.realize(a.grad)
np.testing.assert_allclose(a.grad.numpy(), [2., 2., 2., 2.])
def test_custom_kernel_precompile_multidevice(self):
# a custom_kernel output placeholder (invalids) under multi-device @function(precompile=True) must return the
# kernel's computed result. read it back through .numpy() so the cross-device gather reads the output buffer
devs = ("CPU:0", "CPU:1")
def double_kernel(C:UOp, A:UOp) -> UOp:
C, A = C.flatten(), A.flatten()
i = UOp.range(A.numel(), 0)
return C[i].store(A[i] * 2.0).end(i).sink(arg=KernelInfo(name="double_kernel"))
def double_grad(d_c:UOp, call:UOp): return (None, (Tensor(d_c) * 2.0).uop)
a = Tensor.full((4, 4), 7.0).contiguous().shard(devs, axis=0).realize()
@function(precompile=True, precompile_backward=True)
def f(a:Tensor):
c = Tensor(Tensor.invalids(a.shape[0]//len(devs), a.shape[1], dtype=a.dtype, device=devs).uop.unshard(0), device=devs)
return Tensor.custom_kernel(c, a, fxn=double_kernel, grad_fxn=double_grad)[0]
np.testing.assert_allclose(f(a).numpy(), 14.0)
# g is f with empty output instead of invalids
@function(precompile=True, allow_implicit=True)
def g(a:Tensor):
c = Tensor(Tensor.empty(a.shape[0]//len(devs), a.shape[1], dtype=a.dtype, device=devs).uop.unshard(0), device=devs)
return Tensor.custom_kernel(c, a, fxn=double_kernel, grad_fxn=double_grad)[0]
np.testing.assert_allclose(g(a).numpy(), 14.0)
def test_custom_kernel_inplace_output_is_implicit(self):
# a custom_kernel output the kernel also READS (in-place add) is not write-only, so it must be captured as an input
def inplace_add(C:UOp, A:UOp) -> UOp:
i = UOp.range(A.shape[0], 0)
return C[i].store(C[i].load() + A[i]).end(i).sink(arg=KernelInfo(name="inplace_add"))
@function(precompile=True, allow_implicit=False)
def f(a:Tensor): return Tensor.custom_kernel(Tensor.empty(*a.shape, dtype=a.dtype, device=a.device), a, fxn=inplace_add)[0]
with self.assertRaisesRegex(RuntimeError, "implicit buffer"): f(Tensor([1., 2., 3., 4.]).contiguous().realize())
def test_custom_kernel_write_only_persistent_output_is_implicit(self):
# a write-only custom_kernel output that is a realized buffer must be captured
def write(C:UOp, A:UOp) -> UOp:
i = UOp.range(A.shape[0], 0)
return C[i].store(A[i] * 2.0).end(i).sink(arg=KernelInfo(name="write"))
state = Tensor([100., 200., 300., 400.], device="CPU").contiguous().realize()
@function(precompile=True, allow_implicit=True)
def f(a:Tensor): return Tensor.custom_kernel(state, a, fxn=write)[0]
f(Tensor([1., 2., 3., 4.], device="CPU").contiguous().realize()).realize()
np.testing.assert_allclose(state.numpy(), [2., 4., 6., 8.])
def test_custom_kernel_program_invalids_not_captured(self):
# llama FP8 kernels are PROGRAM with bare-buffer sinks (no analyzable stores), so the invalids scratch
# still must not be captured as an input -- else it is read before the kernel writes it
src = "void k(float* restrict data0, float* restrict data1) { for (int i=0;i<4;i++) data0[i]=data1[i]*2.0f; }"
lib = Device["CPU"].compiler.compile(src)
def prog(C:UOp, A:UOp) -> UOp:
sink = UOp.sink(C.base, A.base, arg=KernelInfo(name="k"))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=(*sink.src, sink)),
UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=lib)),
arg=ProgramInfo(name="k", global_size=(1, 1, 1), local_size=(1, 1, 1), globals=(0, 1)))
@function(precompile=True)
def f(a:Tensor):
c = Tensor.invalids(*a.shape, dtype=a.dtype, device=a.device)
return Tensor.custom_kernel(c, a, fxn=prog)[0]
a = Tensor([1., 2., 3., 4.], device="CPU").contiguous().realize()
np.testing.assert_allclose(f(a).numpy(), [2., 4., 6., 8.])
def test_invalid_store_into_realized_buffer_is_captured(self):
# only fresh invalids() scratch is skipped; a realized buffer is a real input even if an Invalid store
# writes into part of it (its other elements must be preserved), so it is still captured
state = Tensor([10., 20., 30., 40.], device="CPU").contiguous().realize()
@function(precompile=True, allow_implicit=True)
def f(a:Tensor):
after = state.uop.after(state.uop.shrink(((0, 2),)).store(Invalid))
return Tensor(after).contiguous() + a
out = f(Tensor([1., 1., 1., 1.], device="CPU").contiguous().realize())
np.testing.assert_allclose(out.numpy(), [11., 21., 31., 41.])
def test_custom_kernel_precompile_further_compute(self, multi=False, kernel_count:int=2):
devs = ("CPU:0", "CPU:1")
def my_kernel(C:UOp, A:UOp) -> UOp:
C, A = C.flatten(), A.flatten()
i = UOp.range(A.numel(), 0)
return C[i].store(A[i] * 2.0).end(i).sink(arg=KernelInfo(name="my_kernel"))
@function(precompile=True)
def f(a:Tensor):
c = Tensor.invalids(*a.uop.shard_shape, dtype=a.dtype, device=a.device)
if multi: c = Tensor(c.uop.unshard(a.uop.axis), device=a.device)
c = Tensor.custom_kernel(c, a, fxn=my_kernel)[0]
return c + 1
a = Tensor([1., 2., 3., 4.]).contiguous()
if multi: a = a.shard(devs, axis=0)
a.realize()
out = f(a)
GlobalCounters.reset()
out.realize()
assert_kernel_count(kernel_count)
np.testing.assert_allclose(out.numpy(), [3., 5., 7., 9.])
def test_custom_kernel_precompile_further_compute_multi(self): self.test_custom_kernel_precompile_further_compute(multi=True, kernel_count=4)
class TestFunctionGrad(unittest.TestCase):
def test_function_grad_ops(self, precompile=False, precompile_backward=False):
N = 64
x = Tensor.ones(N,N).contiguous()
w1 = Tensor.ones(N,N).contiguous()
w2 = Tensor.ones(N,N).contiguous()
w3 = Tensor.ones(N,N).contiguous()
ref = Tensor.ones(N,N).contiguous()
Tensor.realize(x, w1, w2, w3, ref)
@function(precompile=precompile, precompile_backward=precompile_backward)
def f(x, w1, w2, w3) -> tuple[Tensor, ...]:
p1 = x@w1
p2 = p1@w2
p3 = p2@w3
return p1, p2, p3, p3.contiguous()
ret = f(x, w1, w2, w3)[-1]
loss = (ret-ref).square().mean().backward()
print("RESET")
GlobalCounters.reset()
loss.realize(w1.grad, w2.grad, w3.grad)
print(GlobalCounters.global_ops, GlobalCounters.global_mem)
self.assertLessEqual(GlobalCounters.global_ops, 5000000)
def test_function_grad_ops_precompile(self): self.test_function_grad_ops(precompile=True)
def test_function_grad_ops_precompile_backward(self):
self.test_function_grad_ops(precompile=True, precompile_backward=True)
if __name__ == '__main__':
unittest.main()

21
artifacts/package_sources/tinygrad/test/unit/test_getitem_ops.py Normal file
View File

@@ -0,0 +1,21 @@
import unittest
import numpy as np
from tinygrad import Tensor, GlobalCounters
class TestGetitemOps(unittest.TestCase):
def test_two_tensor_indices(self):
# linear indexing is O(idx_size), one-hot masks is O(idx_size * src_size)
src_np = np.random.rand(10, 100, 200).astype(np.float32)
idx1_np, idx2_np = np.random.randint(0, 100, (50, 60), dtype=np.int32), np.random.randint(0, 200, (50, 60), dtype=np.int32)
src, idx1, idx2 = Tensor(src_np), Tensor(idx1_np), Tensor(idx2_np)
# O(50*60) = 3K vs O(50*60*100*200) = 60M
GlobalCounters.reset()
np.testing.assert_allclose(src_np[0, idx1_np, idx2_np], src[0, idx1, idx2].numpy())
self.assertLess(GlobalCounters.global_ops, 100_000)
# consecutive indices not starting from dim 0: O(10*50*60) = 30K vs O(10*50*60*100*200) = 600M
GlobalCounters.reset()
np.testing.assert_allclose(src_np[:, idx1_np, idx2_np], src[:, idx1, idx2].numpy())
self.assertLess(GlobalCounters.global_ops, 1_000_000)
if __name__ == '__main__':
unittest.main()

260
artifacts/package_sources/tinygrad/test/unit/test_gguf.py Normal file
View File

@@ -0,0 +1,260 @@
import os, struct, unittest, tempfile, pathlib, sys
from tinygrad import dtypes, Tensor, fetch, Device
from tinygrad.helpers import disable_gc
from tinygrad.llm.gguf import _ggml_iq_grid, ggml_data_to_tensor, gguf_load
from tinygrad.runtime.autogen import ggml_common as _ggml
import numpy as np
from gguf import GGUFReader, GGUFValueType, GGMLQuantizationType, GGML_QUANT_SIZES, dequantize, quantize
from gguf.quants import IQ2_S, IQ3_S, IQ3_XXS
ggml_test_block_count = 4
supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes()
class TestGGUFTables(unittest.TestCase):
def test_iq2_s_grid_matches_gguf_py(self):
IQ2_S.init_grid()
grid = _ggml_iq_grid(Device.DEFAULT, _ggml.iq2s_grid, (1024, 8)).numpy()
np.testing.assert_equal(grid, IQ2_S.grid.reshape(1024, 8))
def test_iq3_xxs_grid_matches_gguf_py(self):
IQ3_XXS.init_grid()
grid = _ggml_iq_grid(Device.DEFAULT, _ggml.iq3xxs_grid, (256, 4)).numpy()
np.testing.assert_equal(grid, IQ3_XXS.grid.reshape(256, 4))
def test_iq3_s_grid_matches_gguf_py(self):
IQ3_S.init_grid()
grid = _ggml_iq_grid(Device.DEFAULT, _ggml.iq3s_grid, (512, 4)).numpy()
np.testing.assert_equal(grid, IQ3_S.grid.reshape(512, 4))
@unittest.skipUnless(dtypes.uint8 in supported_dtypes and dtypes.half in supported_dtypes, "Backend must support uint8 and half")
class TestGGUF(unittest.TestCase):
def test_load_tinyllama_q8_0(self): self._test_gguf_load("https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories15M-q8_0.gguf?download=true")
def test_load_tinyllama_q4_0(self): self._test_gguf_load("https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories15M-q4_0.gguf?download=true")
def test_load_gpt2_q4_1(self): self._test_gguf_load("https://huggingface.co/PrunaAI/gpt2-GGUF-smashed/resolve/main/gpt2.Q4_1.gguf?download=true")
def test_load_sample_q6_k(self): self._test_gguf_load("https://huggingface.co/Isotr0py/test-gguf-sample/resolve/main/Quant_Q6_K_1024.gguf?download=true")
def test_dequantization_q8_0_hardcoded(self):
# Q8_0: 2 bytes float16 scale + 32 bytes int8 values, dequant = scale * values
block = np.frombuffer(np.float16(2.0).tobytes() + np.arange(1, 33, dtype=np.int8).tobytes(), dtype=np.uint8).copy()
expected = np.arange(1, 33, dtype=np.float32) * 2.0
np.testing.assert_equal(ggml_data_to_tensor(Tensor(block), 32, GGMLQuantizationType.Q8_0.value).numpy().flatten(), expected)
def test_dequantization_mxfp4_hardcoded(self):
# MXFP4: 1 byte shared exponent E + 16 packed bytes (32 x 4-bit values)
# nibble: bit3=sign, bit2:1=exp, bit0=mant; E=128 gives scale=1.0
# codes 0-7 = [0, 1, 2, 3, 4, 6, 8, 12], codes 8-15 are their negatives
block = np.array([0x80] + list(range(16)), dtype=np.uint8) # E=128, nibbles 0-15 in low, zeros in high
expected = np.array([0., 1., 2., 3., 4., 6., 8., 12., -0., -1., -2., -3., -4., -6., -8., -12.] + [0.]*16, dtype=np.float32)
np.testing.assert_equal(ggml_data_to_tensor(Tensor(block), 32, GGMLQuantizationType.MXFP4.value).numpy().flatten(), expected)
def test_dequantization_q4_0(self): self._test_dequantization(GGMLQuantizationType.Q4_0)
def test_dequantization_q4_1(self): self._test_dequantization(GGMLQuantizationType.Q4_1)
def test_dequantization_q5_0(self): self._test_dequantization(GGMLQuantizationType.Q5_0)
def test_dequantization_q5_1(self): self._test_dequantization(GGMLQuantizationType.Q5_1)
def test_dequantization_q8_0(self): self._test_dequantization(GGMLQuantizationType.Q8_0)
def test_dequantization_q4_k(self): self._test_dequantization(GGMLQuantizationType.Q4_K)
def test_dequantization_q5_k(self): self._test_dequantization(GGMLQuantizationType.Q5_K)
def test_dequantization_q6_k(self): self._test_dequantization(GGMLQuantizationType.Q6_K)
def test_dequantization_iq3_xxs(self): self._test_dequantization(GGMLQuantizationType.IQ3_XXS)
def test_dequantization_iq3_s(self): self._test_dequantization(GGMLQuantizationType.IQ3_S)
def test_dequantization_iq2_s(self): self._test_dequantization(GGMLQuantizationType.IQ2_S)
def test_dequantization_iq4_xs(self): self._test_dequantization(GGMLQuantizationType.IQ4_XS)
def test_dequantization_mxfp4(self): self._test_dequantization(GGMLQuantizationType.MXFP4)
@unittest.skipUnless(dtypes.bfloat16 in supported_dtypes, "Backend must support bfloat16")
def test_dequantization_bf16(self): self._test_dequantization(GGMLQuantizationType.BF16)
def test_dequantization_mxfp4_old(self):
def encode(nibbles, E):
packed = [(low & 0xF) | ((high & 0xF) << 4) for low, high in zip(nibbles[:16], nibbles[16:])]
return np.array([E] + packed, dtype=np.uint8)
def decode(code, E):
sign = -1.0 if (code & 0b1000) else 1.0
exp = (code >> 1) & 0b11
mant = code & 0b1
val = 2 * ((1.0 + 0.5 * mant) * np.exp2(exp - 1) if exp else 0.5 * mant)
scale = np.exp2(E - 128) if E >= 2 else np.exp2(-127 if E == 1 else -128)
return sign * val * scale
blocks, expected = [], []
rng = np.random.default_rng(42)
for _ in range(4):
E = rng.integers(0, 256)
codes = rng.integers(0, 16, size=32, dtype=np.uint8)
blocks.append(encode(codes, E))
expected.extend(decode(c, E) for c in codes)
tensor = Tensor(np.concatenate(blocks))
out = ggml_data_to_tensor(tensor, len(expected), GGMLQuantizationType.MXFP4.value)
np.testing.assert_equal(out.numpy(), expected)
def test_dequantization_mxfp4_block(self):
# https://gist.github.com/Ananta-Ranganathan/3317b6ed51a3b033e9c2564fafb4e043
# used the above script to download the first block of blk.0.attn_k_b.weight from
# https://huggingface.co/unsloth/GLM-4.7-Flash-GGUF/blob/main/GLM-4.7-Flash-MXFP4_MOE.gguf
# and compute the canonical expected dequantized output with the GGUF PY implementation
block = np.array([0x7a, 0x29, 0xab, 0x61, 0x10, 0x21, 0x02, 0x4a,
0x15, 0xca, 0x05, 0x01, 0x9b, 0x39, 0x0b, 0x0b, 0x1c], dtype=np.uint8)
expected = np.array([-0.01562500, -0.04687500, 0.01562500, 0.00000000,
0.01562500, 0.03125000, -0.03125000, 0.09375000,
-0.03125000, 0.09375000, 0.01562500, -0.04687500,
-0.01562500, -0.04687500, -0.04687500, -0.06250000,
0.03125000, -0.03125000, 0.12500000, 0.01562500,
0.03125000, 0.00000000, 0.06250000, 0.01562500,
-0.06250000, 0.00000000, 0.00000000, -0.01562500,
0.04687500, 0.00000000, 0.00000000, 0.01562500], dtype=np.float32)
out = ggml_data_to_tensor(Tensor(block), 32, GGMLQuantizationType.MXFP4.value)
np.testing.assert_equal(out.numpy(), expected)
def test_dequantization_q1_0(self):
# Q1_0: 2 bytes fp16 scale + 16 bytes (128 1-bit values)
block = np.frombuffer(np.float16(2.0).tobytes() + np.packbits(np.random.choice([0, 1], size=128)).tobytes(), dtype=np.uint8).copy()
expected = np.float16(2.0) * (np.unpackbits(block[2:], bitorder="little").astype(np.int8) * 2 - 1)
# TODO: replace 41 with GGMLQuantizationType.Q1_0.value on next gguf-py release
np.testing.assert_equal(ggml_data_to_tensor(Tensor(block), 128, 41).numpy().flatten(), expected)
def test_expected_failure_unknown_type(self):
with self.assertRaises(ValueError):
ggml_data_to_tensor(Tensor.empty(512, dtype=dtypes.uint8), 256, 1337)
@staticmethod
def _build_gguf(tensors, kvs):
# [header] [kv_data] [tensor_infos] [padding] [tensor_data_blob]
buf = bytearray()
# Header: magic "GGUF" + version=3 + n_tensors + n_kv
buf += struct.pack("<4siqq", b"GGUF", 3, len(tensors), len(kvs))
# KV entries: [key_len: uint64][key bytes][type: int32][value]
for k, v in kvs:
kb = k.encode()
if isinstance(v, str): buf += struct.pack("<Q", len(kb)) + kb + struct.pack("<i", 8) + struct.pack("<Q", len(v)) + v.encode()
else: buf += struct.pack("<Q", len(kb)) + kb + struct.pack("<i", 4) + struct.pack("<I", v)
data_off = 0
# Tensor infos: [name_len][name][ndims][dims reversed][qtype][offset_into_data_blob]
for name, dims, qtype, data in tensors:
nb = name.encode()
buf += struct.pack("<Q", len(nb)) + nb + struct.pack("<I", len(dims))
for d in reversed(dims): buf += struct.pack("<Q", d)
buf += struct.pack("<i", qtype) + struct.pack("<Q", data_off)
data_off += len(data)
buf += b"\x00" * ((32 - len(buf) % 32) % 32)
for _, _, _, data in tensors: buf += data
return bytes(buf)
def test_multi_part_load(self):
with tempfile.TemporaryDirectory() as d:
d = pathlib.Path(d)
a, b = np.array([1.0, 2.0, 3.0, 4.0], dtype=np.float32), np.array([5.0, 6.0], dtype=np.float32)
(d / "test-00001-of-00002.gguf").write_bytes(self._build_gguf([("a", (4,), 0, a.tobytes())], [("split.count", 2), ("split.no", 0)]))
(d / "test-00002-of-00002.gguf").write_bytes(self._build_gguf([("b", (2,), 0, b.tobytes())], [("split.count", 2), ("split.no", 1)]))
kv, ts = gguf_load(d / "test-00001-of-00002.gguf")
self.assertEqual(kv["split.count"], 2)
np.testing.assert_equal(ts["a"].numpy(), a)
np.testing.assert_equal(ts["b"].numpy(), b)
# missing part 2
(d / "test-00002-of-00002.gguf").unlink()
with self.assertRaises(FileNotFoundError):
gguf_load(d / "test-00001-of-00002.gguf")
def _test_dequantization(self, qtype: GGMLQuantizationType):
block_size, type_size = GGML_QUANT_SIZES[qtype]
n_el, n_bytes = ggml_test_block_count * block_size, ggml_test_block_count * type_size
try:
q_data = quantize((np.random.random((n_el,)).astype(np.float32) * 100 - 50), qtype)
except NotImplementedError:
q_data = np.random.default_rng(42).integers(0, 256, size=n_bytes, dtype=np.uint8)
ref = dequantize(q_data, qtype)
q_tensor = Tensor(q_data)
dq_tensor = ggml_data_to_tensor(q_tensor, n_el, qtype.value).reshape(n_el)
np.testing.assert_equal(dq_tensor.numpy(), ref)
def _test_gguf_load(self, url: str):
fp = fetch(url)
model_size = os.stat(fp).st_size
gguf_tensor = Tensor.empty(model_size, dtype=dtypes.uint8, device=f"disk:{fp}").to(Device.DEFAULT)
kv_data, tensors = gguf_load(gguf_tensor)
reader = GGUFReader(fp)
for rt in reader.tensors:
ref = dequantize(rt.data, rt.tensor_type)
np.testing.assert_equal(tensors[rt.name].numpy(), ref.reshape(tensors[rt.name].shape))
for k, f in reader.fields.items():
if k.startswith("GGUF."): continue # skip file header keys (version, tensor_count, kv_count)
def read_val(i, parts=f.parts, is_str=(f.types[-1] == GGUFValueType.STRING)):
return bytes(parts[i]).decode("utf-8") if is_str else parts[i][0].item()
if f.types[0] == GGUFValueType.ARRAY:
self.assertEqual(kv_data[k], [read_val(i) for i in f.data])
else:
self.assertEqual(kv_data[k], read_val(-1))
class TestGGUFGEMV(unittest.TestCase):
def _test_gguf_gemv(self, qtype: GGMLQuantizationType):
block_size, type_size = GGML_QUANT_SIZES[qtype]
rows, cols = (1024, 512) if qtype == GGMLQuantizationType.BF16 else (8192, 2048)
n_blocks = rows * cols // block_size
rng = np.random.default_rng(42)
if qtype == GGMLQuantizationType.BF16:
q_data = (rng.standard_normal(rows * cols).astype(np.float32).view(np.uint32) >> 16).astype(np.uint16).view(np.uint8)
else:
# generate random quantized blocks with valid fp16 scale fields (random bytes can produce NaN scales)
q_data = rng.integers(0, 256, size=n_blocks * type_size, dtype=np.uint8).reshape(n_blocks, type_size)
scales = np.float16(rng.standard_normal(n_blocks * 4)).view(np.uint8).reshape(n_blocks, -1)
if qtype in (GGMLQuantizationType.Q5_0, GGMLQuantizationType.Q8_0,
GGMLQuantizationType.IQ3_XXS,
GGMLQuantizationType.IQ2_S,
GGMLQuantizationType.IQ3_S, GGMLQuantizationType.IQ4_XS): q_data[:, :2] = scales[:, :2] # d at offset 0
elif qtype in (GGMLQuantizationType.Q5_1, GGMLQuantizationType.Q4_K, GGMLQuantizationType.Q5_K):
q_data[:, :4] = scales[:, :4] # d, m/dmin at offset 0
elif qtype == GGMLQuantizationType.Q6_K: q_data[:, -2:] = scales[:, :2] # d at end
elif qtype == GGMLQuantizationType.MXFP4: q_data[:, 0] = rng.integers(120, 136, size=n_blocks, dtype=np.uint8) # constrain byte0
q_data = q_data.flatten()
ref = dequantize(q_data, qtype).reshape(rows, cols)
# build a minimal gguf in memory: header + 1 tensor info + aligned data
buf = bytearray()
buf += struct.pack("<4siqq", b"GGUF", 3, 1, 0) # magic, version, n_tensors, n_kv
buf += struct.pack("<Q", 6) + b"weight" # tensor name
buf += struct.pack("<I", 2) # ndims
buf += struct.pack("<QQ", cols, rows) # dims (gguf stores reversed)
buf += struct.pack("<i", qtype.value)
buf += struct.pack("<Q", 0) # offset
buf += b"\x00" * ((32 - len(buf) % 32) % 32) # pad to alignment=32
buf += q_data.tobytes()
_, tensors = gguf_load(Tensor(np.frombuffer(buf, dtype=np.uint8)).to(None))
x = rng.standard_normal(cols).astype(np.float32)
with np.errstate(all='ignore'):
np.testing.assert_allclose((tensors["weight"] @ Tensor(x)).numpy(), ref @ x, atol=1e-2, rtol=1e-2)
if qtype == GGMLQuantizationType.BF16 or dtypes.half in supported_dtypes: np.testing.assert_equal(tensors["weight"].numpy(), ref)
assert np.isfinite(ref).all() and np.isfinite(tensors["weight"].numpy()).all(), f"{qtype.name} has NaN/Inf"
def test_gguf_gemv_q8_0(self): self._test_gguf_gemv(GGMLQuantizationType.Q8_0)
def test_gguf_gemv_q5_0(self): self._test_gguf_gemv(GGMLQuantizationType.Q5_0)
def test_gguf_gemv_q5_1(self): self._test_gguf_gemv(GGMLQuantizationType.Q5_1)
def test_gguf_gemv_q4_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q4_K)
def test_gguf_gemv_q5_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q5_K)
def test_gguf_gemv_q6_k(self): self._test_gguf_gemv(GGMLQuantizationType.Q6_K)
def test_gguf_gemv_iq3_xxs(self): self._test_gguf_gemv(GGMLQuantizationType.IQ3_XXS)
def test_gguf_gemv_iq3_s(self): self._test_gguf_gemv(GGMLQuantizationType.IQ3_S)
def test_gguf_gemv_iq2_s(self): self._test_gguf_gemv(GGMLQuantizationType.IQ2_S)
def test_gguf_gemv_iq4_xs(self): self._test_gguf_gemv(GGMLQuantizationType.IQ4_XS)
def test_gguf_gemv_mxfp4(self): self._test_gguf_gemv(GGMLQuantizationType.MXFP4)
@unittest.skipUnless(dtypes.bfloat16 in supported_dtypes, "Backend must support bfloat16")
def test_gguf_gemv_bf16(self): self._test_gguf_gemv(GGMLQuantizationType.BF16)
class TestGGUFGC(unittest.TestCase):
def test_gguf_load_no_tensor_leak(self):
"""gguf_load must not retain references to the input tensor after returning."""
fp = fetch("https://huggingface.co/ggml-org/models/resolve/main/tinyllamas/stories15M-q8_0.gguf?download=true")
t = Tensor.empty(os.stat(fp).st_size, dtype=dtypes.uint8, device=f"disk:{fp}").to(Device.DEFAULT).realize()
with disable_gc():
ref_before = sys.getrefcount(t)
kv_data, tensors = gguf_load(t)
self.assertEqual(sys.getrefcount(t), ref_before, "gguf_load leaked a reference to the input tensor")
if __name__ == '__main__':
unittest.main()

210
artifacts/package_sources/tinygrad/test/unit/test_gradient.py Normal file
View File

@@ -0,0 +1,210 @@
import unittest, math
import numpy as np
from tinygrad import Tensor
from tinygrad.dtype import dtypes
from tinygrad.uop.ops import UOp, KernelInfo, Ops
class TestTensorGradient(unittest.TestCase):
def test_example(self):
x = Tensor.eye(3)
y = Tensor([[2.0,0,-2.0]])
z = y.matmul(x).sum()
dx, dy = z.gradient(x, y)
self.assertListEqual(dx.tolist(), [[2.0, 2.0, 2.0], [0.0, 0.0, 0.0], [-2.0, -2.0, -2.0]])
self.assertListEqual(dy.tolist(), [[1.0, 1.0, 1.0]])
def test_zero_if_not_used(self):
x = Tensor([1.0, 2.0, 3.0])
w = Tensor.randn((3,))
self.assertListEqual(x.sum().gradient(w)[0].tolist(), [0.0, 0.0, 0.0])
def test_with_custom_gradient(self):
x = Tensor([1.0, 2.0, 3.0])
z = (x * x).sum()
dx = z.gradient(x, gradient=Tensor([3.0]))[0]
self.assertListEqual(dx.tolist(), [6.0, 12.0, 18.0])
def test_broadcast_gradient(self):
x = Tensor([[1.0], [2.0], [3.0]])
y = Tensor([[10.0, 20.0, 30.0, 40.0]])
z = (x + y).sum()
dx, dy = z.gradient(x, y)
self.assertListEqual(dx.tolist(), [[4.0], [4.0], [4.0]])
self.assertListEqual(dy.tolist(), [[3.0, 3.0, 3.0, 3.0]])
def test_non_scalar_output(self):
x = Tensor([1.0, 2.0, 3.0])
z = x * x
with self.assertRaises(AssertionError): z.gradient(x)
dz = Tensor([1.0, 1.0, 1.0])
dx = z.gradient(x, gradient=dz)[0]
self.assertListEqual(dx.tolist(), [2.0, 4.0, 6.0])
def test_cast_before_view(self):
x = Tensor([1.0, 1, 1, 1])
x_reshaped = x.reshape(2,2)
x_casted = x_reshaped.cast(dtypes.float16)
x_casted.mean().gradient(x_reshaped)
def test_non_float_tensor_raise(self):
x = Tensor([1, 2, 3])
with self.assertRaises(RuntimeError): x.sum().gradient(x)
with self.assertRaises(RuntimeError): x.float().sum().gradient(x)
def test_copy_to_device_gradient(self):
t = Tensor([1.0, 2, 3]).realize()
t.to("CPU:1").square().sum().backward()
self.assertEqual(t.grad.device, t.device)
self.assertListEqual(t.grad.tolist(), [2.0, 4.0, 6.0])
def test_multiple_backward(self):
x = Tensor([3.])
(x*2)[0].backward()
np.testing.assert_allclose(x.grad.numpy(), [2.0])
old_grad = x.grad
(x*3)[0].backward()
np.testing.assert_allclose(x.grad.numpy(), [2.0+3.0])
self.assertIs(x.grad, old_grad)
(x*x)[0].backward()
np.testing.assert_allclose(x.grad.numpy(), [2.0+3.0+2*3.0])
self.assertIs(x.grad, old_grad)
def test_gradient_through_clone_from_grad_src(self):
# unlike torch, tinygrad accumulates grad on every tensor in the graph, including non-leaf x
src = Tensor([1.0, 2.0, 3.0, 4.0])
x = src.clone()
(x * 2.0).sum().backward()
np.testing.assert_allclose(src.grad.numpy(), [2.0, 2.0, 2.0, 2.0])
np.testing.assert_allclose(x.grad.numpy(), [2.0, 2.0, 2.0, 2.0])
def test_gradient_through_clone_from_detached_src(self):
base = Tensor([1.0, 2.0, 3.0, 4.0])
x = base.detach().clone()
(x * 2.0).sum().backward()
np.testing.assert_allclose(x.grad.numpy(), [2.0, 2.0, 2.0, 2.0]) # gradient flows through clone
np.testing.assert_allclose(base.grad.numpy(), [0.0, 0.0, 0.0, 0.0]) # ...but detach blocks it from base
def test_setitem_on_grad_used_tensor_raises(self):
x = Tensor([1.0, 2.0, 3.0, 4.0]).realize()
_ = (x * 2.0).sum()
with self.assertRaises(RuntimeError):
x[0] = 99.0
def test_gradient_through_chained_unrealized_setitem(self):
g1 = Tensor.zeros(4).contiguous()
g1[2] = Tensor(1.0)
g2 = Tensor.zeros(5, 4).contiguous()
g2[0] = g1
x = Tensor.randn(4, 4)
np.testing.assert_allclose(x.pad(((1,0),(0,0))).gradient(x, gradient=g2)[0].numpy(), np.zeros((4, 4)))
def test_implicit_broadcast_where_gradient(self):
# WHERE with a bare ()-shape branch: the scalar's gradient counts the positions where it is selected
cond, x, w = Tensor([True, False, True]), Tensor([1.0, 2.0, 3.0]), Tensor(4.0)
dw = Tensor(cond.uop.alu(Ops.WHERE, x.uop, w.uop)).sum().gradient(w)[0]
self.assertEqual(dw.shape, ())
self.assertEqual(dw.item(), 1.0)
dw = Tensor(cond.uop.alu(Ops.WHERE, w.uop, x.uop)).sum().gradient(w)[0]
self.assertEqual(dw.item(), 2.0)
def test_implicit_broadcast_alu_gradient(self):
# MUL with a bare ()-shape src, no EXPAND in the graph
x, w = Tensor([1.0, 2.0, 3.0]), Tensor(2.0)
m = x.uop.alu(Ops.MUL, w.uop)
self.assertIs(m.src[1], w.uop)
dw = Tensor(m).sum().gradient(w)[0]
self.assertEqual(dw.shape, ())
self.assertEqual(dw.item(), 6.0)
def test_implicit_broadcast_intermediate_accumulation(self):
# s is used directly and through an implicit broadcast edge, each edge's gradient reduces to s's shape before they sum
x, p = Tensor([1.0, 2.0, 3.0]), Tensor(0.5)
s = p.sin()
z = Tensor(x.uop.alu(Ops.MUL, s.uop)).sum() + s
dp = z.gradient(p)[0]
self.assertEqual(dp.shape, ())
self.assertAlmostEqual(dp.item(), 7*math.cos(0.5), places=5)
def test_bare_const_skipped_by_backward(self):
Tensor.manual_seed(0)
w = Tensor(1.0)
(Tensor.rand(()) + w).backward()
self.assertIsNone(w.grad)
class TestMultiOutputGradient(unittest.TestCase):
@staticmethod
def addmul_kernel(C:UOp, D:UOp, A:UOp, B:UOp) -> UOp:
C, D, A, B = C.flatten(), D.flatten(), A.flatten(), B.flatten()
i = UOp.range(C.numel(), 0)
store_c = C[i].store(A[i] + B[i])
store_d = D[i].store(A[i] * B[i])
return UOp.group(store_c, store_d).end(i).sink(arg=KernelInfo(name="addmul")).simplify()
@staticmethod
def backward_addmul(grad_c, grad_d, call):
_c, _d, a, b = call.src[1:]
grad_a = (Tensor(grad_c) + Tensor(grad_d) * Tensor(b)).uop
grad_b = (Tensor(grad_c) + Tensor(grad_d) * Tensor(a)).uop
return (None, None, grad_a, grad_b)
def test_custom_kernel_multi_output_backward(self):
a_np, b_np = np.random.randn(4, 4).astype(np.float32), np.random.randn(4, 4).astype(np.float32)
a_ref, b_ref = Tensor(a_np), Tensor(b_np)
((a_ref + b_ref).sum() + (a_ref * b_ref).sum()).backward()
a, b = Tensor(a_np), Tensor(b_np)
Tensor.realize(a, b)
c, d, _, _ = Tensor.custom_kernel(Tensor.empty(4, 4), Tensor.empty(4, 4), a, b, fxn=self.addmul_kernel, grad_fxn=self.backward_addmul)
(c.sum() + d.sum()).backward()
np.testing.assert_allclose(a.grad.numpy(), a_ref.grad.numpy(), rtol=1e-5)
np.testing.assert_allclose(b.grad.numpy(), b_ref.grad.numpy(), rtol=1e-5)
def test_custom_kernel_multi_output_backward_interacting(self):
a_np, b_np = np.random.randn(4, 4).astype(np.float32), np.random.randn(4, 4).astype(np.float32)
a_ref, b_ref = Tensor(a_np), Tensor(b_np)
((a_ref + b_ref) * (a_ref * b_ref)).sum().backward()
a, b = Tensor(a_np), Tensor(b_np)
Tensor.realize(a, b)
c, d, _, _ = Tensor.custom_kernel(Tensor.empty(4, 4), Tensor.empty(4, 4), a, b, fxn=self.addmul_kernel, grad_fxn=self.backward_addmul)
(c * d).sum().backward()
np.testing.assert_allclose(a.grad.numpy(), a_ref.grad.numpy(), rtol=1e-5, atol=1e-7)
np.testing.assert_allclose(b.grad.numpy(), b_ref.grad.numpy(), rtol=1e-5, atol=1e-7)
def test_custom_kernel_three_output_backward(self):
def addmulsub_kernel(C:UOp, D:UOp, E:UOp, A:UOp, B:UOp) -> UOp:
C, D, E, A, B = C.flatten(), D.flatten(), E.flatten(), A.flatten(), B.flatten()
i = UOp.range(C.numel(), 0)
store_c = C[i].store(A[i] + B[i])
store_d = D[i].store(A[i] * B[i])
store_e = E[i].store(A[i] - B[i])
return UOp.group(store_c, store_d, store_e).end(i).sink(arg=KernelInfo(name="addmulsub")).simplify()
def backward_addmulsub(grad_c, grad_d, grad_e, call):
_c, _d, _e, a, b = call.src[1:]
grad_a = (Tensor(grad_c) + Tensor(grad_d) * Tensor(b) + Tensor(grad_e)).uop
grad_b = (Tensor(grad_c) + Tensor(grad_d) * Tensor(a) - Tensor(grad_e)).uop
return (None, None, None, grad_a, grad_b)
a_np, b_np = np.random.randn(4, 4).astype(np.float32), np.random.randn(4, 4).astype(np.float32)
a_ref, b_ref = Tensor(a_np), Tensor(b_np)
((a_ref + b_ref).sum() + (a_ref * b_ref).sum() + (a_ref - b_ref).sum()).backward()
a, b = Tensor(a_np), Tensor(b_np)
Tensor.realize(a, b)
c, d, e, _, _ = Tensor.custom_kernel(Tensor.empty(4, 4), Tensor.empty(4, 4), Tensor.empty(4, 4), a, b,
fxn=addmulsub_kernel, grad_fxn=backward_addmulsub)
(c.sum() + d.sum() + e.sum()).backward()
np.testing.assert_allclose(a.grad.numpy(), a_ref.grad.numpy(), atol=1e-6, rtol=1e-5)
np.testing.assert_allclose(b.grad.numpy(), b_ref.grad.numpy(), atol=1e-6, rtol=1e-5)
class TestViewGradient(unittest.TestCase):
def test_expand(self):
x = Tensor.randn(5,2)
a = Tensor([3.])
aex = a.expand(10)
(aex.reshape(5,2) * x).sum().backward()
np.testing.assert_allclose(aex.grad.numpy(), x.reshape(10).numpy())
with self.assertRaises(AssertionError):
np.testing.assert_allclose(aex.grad.numpy(), a.grad.expand(10).numpy())
if __name__ == '__main__':
unittest.main()

109
artifacts/package_sources/tinygrad/test/unit/test_hashing.py Normal file
View File

@@ -0,0 +1,109 @@
from typing_extensions import Callable
import hashlib, random, unittest
from tinygrad import Tensor, Device, dtypes
from tinygrad.helpers import DEV
from test.helpers import slow
from tinygrad.uop.ops import UOp
from tinygrad.engine.jit import TinyJit
supported_dtypes = Device[Device.DEFAULT].renderer.supported_dtypes()
@unittest.skipUnless(dtypes.uint8 in supported_dtypes and dtypes.uint64 in supported_dtypes, "Device must support uint8 and uint64")
@unittest.skipIf(DEV.interface.startswith("MOCK") and Device.DEFAULT == "NV", "crashes in NV CI")
class TestHashing(unittest.TestCase):
def _python_hash_1mb(self, data:bytes):
chunks = [data[i:i+4096] for i in range(0, len(data), 4096)]
chunk_hashes = [hashlib.shake_128(chunk).digest(16) for chunk in chunks]
return hashlib.shake_128(b''.join(chunk_hashes)).digest(16)
@unittest.skip("very slow")
def test_abc(self):
expected = self._python_hash_1mb(b"abc" + b"\x00" * (2**20 - 3))
out = Tensor(b"abc").hash()
self.assertEqual(bytes(out.data()), expected)
@unittest.skipUnless(dtypes.uint8 in supported_dtypes and dtypes.uint64 in supported_dtypes, "Device must support uint8 and uint64")
@unittest.skipIf(DEV.interface.startswith("MOCK") and Device.DEFAULT == "NV", "crashes in NV CI")
class TestKeccak(unittest.TestCase):
def setUp(self) -> None: random.seed(1337)
def test_shape_keeping(self):
s = (1, 2, 3, 4)
for i in range(len(s)):
out_shape = Tensor.randint(*s[i:], high=255, dtype=dtypes.uint8).keccak().shape
self.assertTupleEqual(s[i:-1], out_shape[:-1])
@slow
def test_sha3_224(self): self._test_preset("sha3_224", [143, 144])
@slow
def test_sha3_256(self): self._test_preset("sha3_256", [135, 136])
@slow
def test_shake_128(self): self._test_preset("shake_128", [167, 168], lambda d: hashlib.shake_128(d).digest(16))
def _test_preset(self, name: str, special_sizes: list[int], hasher: Callable[[bytes], bytes] | None = None):
def default_hasher(d: bytes) -> bytes: return getattr(hashlib, name)(d).digest()
if hasher is None: hasher = default_hasher
for n in (special_sizes + [special_sizes[0] - 1]):
a, b = random.randbytes(n), random.randbytes(n)
ha_ref, hb_ref = hasher(a), hasher(b)
tres = Tensor.stack(*(Tensor(d) for d in (a, b))).keccak(name)
ha, hb = bytes(tres[0].data()), bytes(tres[1].data())
self.assertEqual(ha_ref, ha)
self.assertEqual(ha_ref, bytes(Tensor(a).keccak(name).data()))
self.assertEqual(hb_ref, hb)
def test_referenced(self):
# https://www.di-mgt.com.au/sha_testvectors.html
self.assertEqual(bytes(Tensor(b"abc").keccak().tolist()),
bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
@slow
def test_long(self):
data = b"\x00" * 4
self.assertEqual(bytes(Tensor(data).keccak("shake_128").tolist()), hashlib.shake_128(data).digest(16))
data = b"\x00" * 1000
self.assertEqual(bytes(Tensor(data).keccak("shake_128").tolist()), hashlib.shake_128(data).digest(16))
def test_variable_bs(self):
data = Tensor([b"abc", b"abc", b"def"], dtype=dtypes.uint8).repeat(2048, 1)
bs = UOp.variable("bs", 1, 4096).bind(3)
out = data.shrink_to(bs, data.shape[-1]).keccak().shrink_to(3, 32).realize()
self.assertEqual(bytes(out[0].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
self.assertEqual(bytes(out[1].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
self.assertEqual(bytes(out[2].tolist()), bytearray.fromhex("8e0d8f672252acb0 ffc5093db8653b18 1513bf9a2097e737 b4f73533dcaf46df"))
@slow
def test_variable_bs_jit(self):
def f(data):
return data.keccak()
jit_f = TinyJit(f)
data = Tensor([b"abc", b"abc", b"abc"], dtype=dtypes.uint8).repeat(2048, 1)
# initialize jit
for _ in range(3):
bs = UOp.variable("bs", 1, 4096).bind(4096)
_ = jit_f(data.shrink_to(bs, data.shape[-1]))
bs = UOp.variable("bs", 1, 4096).bind(1)
out = jit_f(data.shrink_to(bs, data.shape[-1])).shrink_to(1, 32)
self.assertEqual(bytes(out[0].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
bs = UOp.variable("bs", 1, 4096).bind(2)
out = jit_f(data.shrink_to(bs, data.shape[-1])).shrink_to(2, 32)
self.assertEqual(bytes(out[0].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
self.assertEqual(bytes(out[1].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
bs = UOp.variable("bs", 1, 4096).bind(3)
data = Tensor([b"abc", b"abc", b"def"], dtype=dtypes.uint8).repeat(2048, 1)
out = jit_f(data.shrink_to(bs, data.shape[-1])).shrink_to(3, 32)
self.assertEqual(bytes(out[0].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
self.assertEqual(bytes(out[1].tolist()), bytearray.fromhex("3a985da74fe225b2 045c172d6bd390bd 855f086e3e9d525b 46bfe24511431532"))
self.assertEqual(bytes(out[2].tolist()), bytearray.fromhex("8e0d8f672252acb0 ffc5093db8653b18 1513bf9a2097e737 b4f73533dcaf46df"))
if __name__ == "__main__":
unittest.main()

44
artifacts/package_sources/tinygrad/test/unit/test_hcq_graph.py Normal file
View File

@@ -0,0 +1,44 @@
import unittest
from tinygrad import Device, Tensor
from tinygrad.engine.jit import TinyJit
from tinygrad.uop.ops import UOp, Ops
from tinygrad.dtype import dtypes
from tinygrad.runtime.graph.hcq import HCQGraph
from tinygrad.runtime.support.hcq import HCQCompiled
from tinygrad.runtime.support.usb import USBMMIOInterface
from test.mockgpu.usb import MockUSB
@unittest.skipUnless(issubclass(type(Device[Device.DEFAULT]), HCQCompiled), "HCQ device required to run")
class TestHCQUnit(unittest.TestCase):
@unittest.skipIf(Device.DEFAULT == "CPU", "requires non-CPU HCQ device")
def test_supports_uop(self):
d0, cpu_dev = Device[Device.DEFAULT], Device["CPU"]
@TinyJit
def f(inp, inp_cpu):
return (inp + 1.0).contiguous().realize(), (inp_cpu + 1.0).contiguous().realize()
inp, inp_cpu = Tensor.randn(10, 10, device=Device.DEFAULT).realize(), Tensor.randn(10, 10, device="CPU").realize()
for _ in range(5): f(inp, inp_cpu)
# construct minimal CALL UOps for supports_uop (graphs only see PROGRAMs after compile_linear)
gpu_call = UOp(Ops.PROGRAM, src=(UOp.sink(),)).call(UOp.new_buffer(Device.DEFAULT, 1, dtypes.float))
cpu_call = UOp(Ops.PROGRAM, src=(UOp.sink(),)).call(UOp.new_buffer("CPU", 1, dtypes.float))
gpu_devs = [d0]
# local MMIO: GPU works alone and with CPU in batch (cpu_support=True)
assert HCQGraph.supports_uop(gpu_devs, gpu_call) is True
assert HCQGraph.supports_uop(gpu_devs, cpu_call) is True
assert HCQGraph.supports_uop(gpu_devs + [cpu_dev], gpu_call) is True
# USB MMIO: GPU-only still works, but CPU batching must be rejected (cpu_support=False)
orig_view = d0.timeline_signal.base_buf.view
try:
d0.timeline_signal.base_buf.view = USBMMIOInterface(MockUSB(bytearray(256)), 0, 16, fmt='B')
assert HCQGraph.supports_uop(gpu_devs, gpu_call) is True
assert HCQGraph.supports_uop(gpu_devs, cpu_call) is False
assert HCQGraph.supports_uop(gpu_devs + [cpu_dev], gpu_call) is False
finally:
d0.timeline_signal.base_buf.view = orig_view
if __name__ == "__main__":
unittest.main()

29
artifacts/package_sources/tinygrad/test/unit/test_helpers.py Normal file
View File

@@ -0,0 +1,29 @@
import unittest, gc
import numpy as np
from tinygrad.helpers import polyN, disable_gc
from tinygrad.tensor import Tensor, is_numpy_ndarray
class TestPolyN(unittest.TestCase):
def test_tensor(self):
np.testing.assert_allclose(polyN(Tensor([1.0, 2.0, 3.0, 4.0]), [1.0, -2.0, 1.0]).numpy(), [0.0, 1.0, 4.0, 9.0])
class TestIsNumpyNdarray(unittest.TestCase):
def test_tensor_numpy(self):
self.assertTrue(is_numpy_ndarray(Tensor([1, 2, 3]).numpy()))
class TestDisableGC(unittest.TestCase):
def test_recursive_decorator(self):
was_enabled = gc.isenabled()
@disable_gc()
def recurse(depth:int):
self.assertFalse(gc.isenabled())
if depth: recurse(depth-1)
self.assertFalse(gc.isenabled())
try:
recurse(2)
self.assertEqual(gc.isenabled(), was_enabled)
finally:
(gc.enable if was_enabled else gc.disable)()
if __name__ == '__main__':
unittest.main()

1393
artifacts/package_sources/tinygrad/test/unit/test_indexing.py Normal file
View File

File diff suppressed because it is too large Load Diff

147
artifacts/package_sources/tinygrad/test/unit/test_invalid_tensor.py Normal file
View File

@@ -0,0 +1,147 @@
import unittest
from tinygrad import Tensor
from tinygrad.device import Buffer
from tinygrad.dtype import Invalid, dtypes
from tinygrad.engine.realize import run_linear
from tinygrad.uop.ops import Ops, UOp
class TestInvalidTensor(unittest.TestCase):
def _invalid_test_helper(self, out, expected):
linear, var_vals = out.linear_with_vars()
buf = out.uop.buffer
buf.allocate()
sentinel = memoryview(bytearray(b'\x42' * buf.nbytes))
buf.copy_from(Buffer("PYTHON", buf.size, buf.dtype, opaque=sentinel))
before = buf.as_memoryview().cast(out.dtype.fmt).tolist()
run_linear(linear, var_vals)
ret = buf.as_memoryview().cast(out.dtype.fmt).tolist()
for i,v in enumerate(expected): self.assertEqual(ret[i], before[i] if v is None else v)
def test_where_x_invalid(self):
mask = Tensor.arange(4) < 2
out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Invalid)
self._invalid_test_helper(out, [1.0, 2.0, None, None])
def test_where_invalid_x(self):
mask = Tensor.arange(4) < 2
out = mask.where(Invalid, Tensor([1.0, 2.0, 3.0, 4.0]))
self._invalid_test_helper(out, [None, None, 3.0, 4.0])
def test_where_invalid_2d(self):
mask = Tensor.arange(6).reshape(2, 3) < 3
vals = Tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
out = mask.where(vals, Invalid)
self._invalid_test_helper(out, [1.0, 2.0, 3.0, None, None, None])
def test_where_invalid_int(self):
mask = Tensor.arange(3) < 2
out = mask.where(Tensor([10, 20, 30]), Invalid)
self._invalid_test_helper(out, [10, 20, None])
def test_where_invalid_add(self):
mask = Tensor.arange(3) < 2
mixed = mask.where(Tensor([10.0, 20.0, 30.0]), Invalid)
out = mixed + Tensor([1.0, 2.0, 3.0])
self._invalid_test_helper(out, [11.0, 22.0, None])
def test_where_invalid_add_left(self):
mask = Tensor.arange(3) < 2
mixed = mask.where(Tensor([10.0, 20.0, 30.0]), Invalid)
out = Tensor([1.0, 2.0, 3.0]) + mixed
self._invalid_test_helper(out, [11.0, 22.0, None])
def test_where_always_true(self):
mask = Tensor.arange(3) < 10
out = mask.where(Tensor([10.0, 20.0, 30.0]), Invalid)
self._invalid_test_helper(out, [10.0, 20.0, 30.0])
def test_where_cast(self):
mask = Tensor.arange(4) < 2
out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Invalid).cast(dtypes.int)
self._invalid_test_helper(out, [1, 2, None, None])
def test_where_compare(self):
mask = Tensor.arange(4) < 2
out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Invalid) > 1
self._invalid_test_helper(out, [False, True, None, None])
def test_where_invalid_condition(self):
a, x = Tensor.arange(4), Tensor([0, 1, 2, 3])
bad = (a < 2).where(a, Invalid)
out = (bad < 1).logical_not().where(x + 10, x + 20)
self._invalid_test_helper(out, [20, 11, None, None])
def test_where_invalid_condition_bare(self):
cond = Tensor.full((4,), Invalid, dtype=dtypes.bool, buffer=False)
out = cond.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor([10.0, 20.0, 30.0, 40.0]))
self._invalid_test_helper(out, [None, None, None, None])
def test_where_unary(self):
mask = Tensor.arange(4) < 2
out = mask.where(Tensor([1.0, 4.0, 9.0, 16.0]), Invalid).sqrt()
self._invalid_test_helper(out, [1.0, 2.0, None, None])
def test_where_where(self):
mask1 = Tensor.arange(4) < 2
mask2 = Tensor.arange(4) > 0
out = mask2.where(mask1.where(Tensor([1.0, 2.0, 3.0, 4.0]), Invalid), Invalid)
self._invalid_test_helper(out, [None, 2.0, None, None])
def test_where_reduce_always_true(self):
mask = Tensor.arange(4) < 9
out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Invalid).sum()
self._invalid_test_helper(out, [10.0])
def test_invalid_unary(self):
mask = Tensor.arange(4) < 2
out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.float, buffer=False).sqrt())
self._invalid_test_helper(out, [1.0, 2.0, None, None])
def test_invalid_binary(self):
mask = Tensor.arange(4) < 2
out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.float, buffer=False) + 2)
self._invalid_test_helper(out, [1.0, 2.0, None, None])
def test_invalid_binary_left(self):
mask = Tensor.arange(4) < 2
out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), 2 + Tensor.full((4,), Invalid, dtype=dtypes.float, buffer=False))
self._invalid_test_helper(out, [1.0, 2.0, None, None])
def test_invalid_reshape(self):
mask = Tensor.arange(4) < 2
out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Invalid).reshape(2,2)
self._invalid_test_helper(out, [1.0, 2.0, None, None])
def test_invalid_cast(self):
mask = Tensor.arange(4) < 2
out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.int, buffer=False).cast(dtypes.float))
self._invalid_test_helper(out, [1.0, 2.0, None, None])
def test_invalid_bitcast(self):
mask = Tensor.arange(4) < 2
out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.int, buffer=False).bitcast(dtypes.float))
self._invalid_test_helper(out, [1.0, 2.0, None, None])
def test_where_bitcast(self):
mask = Tensor.arange(4) < 2
out = mask.where(Tensor([1.0, 2.0, 3.0, 4.0]), Tensor.full((4,), Invalid, dtype=dtypes.int, buffer=False)).bitcast(dtypes.int)
self._invalid_test_helper(out, [0x3f800000, 0x40000000, None, None])
def test_tensor_index(self):
idx = (Tensor.arange(4) < 2).where(Tensor([0, 1, 2, 3]), Invalid)
out = Tensor([1.0, 2.0, 3.0, 4.0])[idx]
self._invalid_test_helper(out, [1.0, 2.0, None, None])
def test_uop_where_keeps_invalid_bare(self):
cond = UOp.const(0) < UOp.const(1)
idx = UOp(Ops.STACK, src=tuple(UOp.const(x) for x in range(3)))
out = cond.where(idx, UOp.invalid())
self.assertIs(cond.op, Ops.CMPLT)
self.assertIs(idx.op, Ops.STACK)
self.assertIs(out.op, Ops.WHERE)
self.assertIs(out.src[2].op, Ops.CONST)
self.assertTrue(out.src[2].is_invalid)
if __name__ == '__main__':
unittest.main()

442
artifacts/package_sources/tinygrad/test/unit/test_jit.py Normal file
View File

@@ -0,0 +1,442 @@
import unittest, numpy as np
from test.helpers import assert_jit_cache_len
from tinygrad import Tensor, TinyJit, Context, UOp, dtypes
from tinygrad.engine.jit import JitError
def _simple_test(add, extract=lambda x: x, N=10):
for _ in range(5):
a = Tensor.randn(N, N)
b = Tensor.randn(N, N)
c = add(a, b)
np.testing.assert_allclose(extract(c).numpy(), a.numpy()+b.numpy(), atol=1e-4, rtol=1e-5)
assert_jit_cache_len(add, 1)
class TestJit(unittest.TestCase):
def test_jitbeam_triggers_beam(self):
from unittest.mock import patch
from tinygrad.helpers import getenv as _getenv
@TinyJit
def add(a, b): return (a+b).realize()
a, b = Tensor.ones(10, 10).contiguous().realize(), Tensor.ones(10, 10).contiguous().realize()
with patch("tinygrad.codegen.opt.search.beam_search", wraps=lambda k,*a,**kw: k) as mock_beam:
add(a, b)
assert mock_beam.call_count == 0
with patch("tinygrad.engine.jit.getenv", side_effect=lambda k, d=0: 1 if k == "JITBEAM" else _getenv(k, d)): add(a, b)
assert mock_beam.call_count == 1
def test_simple_jit_reset(self):
@TinyJit
def add(a, b): return (a+b).realize()
_simple_test(add)
add.reset()
_simple_test(add, N=20)
def test_simple_jit_norealize(self):
@TinyJit
def add(a, b): return (a+b)
_simple_test(add)
def test_simple_jit_norealize_list(self):
@TinyJit
def add(a, b): return [a+b]
_simple_test(add, extract=lambda x: x[0])
def test_simple_jit_norealize_dict(self):
@TinyJit
def add(a, b): return {"billy": a+b}
_simple_test(add, extract=lambda x: x["billy"])
def test_jit_multiple_outputs(self):
@TinyJit
def f(a, b): return (a+b).realize(), (a-b).realize(), (a*b).realize()
for _ in range(5):
a = Tensor.randn(10, 10)
b = Tensor.randn(10, 10)
c, d, e = f(a, b)
np.testing.assert_allclose(c.numpy(), a.numpy()+b.numpy(), atol=1e-4, rtol=1e-5)
np.testing.assert_allclose(d.numpy(), a.numpy()-b.numpy(), atol=1e-4, rtol=1e-5)
np.testing.assert_allclose(e.numpy(), a.numpy()*b.numpy(), atol=1e-4, rtol=1e-5)
assert_jit_cache_len(f, 3)
def test_nothing_jitted(self):
@TinyJit
def add(a, b): return None
with self.assertRaises(JitError):
for _ in range(5):
a = Tensor.randn(10, 10)
b = Tensor.randn(10, 10)
add(a, b)
def test_jit_zero_does_not_jit(self):
@TinyJit
def add(a, b): return (a+b).realize()
with Context(JIT=0):
for i in range(5):
a = Tensor([i])
b = Tensor([i])
c = add(a, b)
np.testing.assert_allclose(c.numpy(), 2*i)
assert_jit_cache_len(add, 0)
def test_jit_not_capturing(self):
@TinyJit
def add(a, b):
Tensor.zeros(4, 4).contiguous().realize() # no-op kernel is captured
return (a+b).realize()
for i in range(5):
a = Tensor([i])
b = Tensor([i])
c = add(a, b)
np.testing.assert_allclose(c.numpy(), 2*i)
assert_jit_cache_len(add, 2)
@TinyJit
def add2(a, b):
with Context(CAPTURING=0): # not captured
Tensor.zeros(4, 4).contiguous().realize()
return (a+b).realize()
for i in range(5):
a = Tensor([i])
b = Tensor([i])
c = add2(a, b)
np.testing.assert_allclose(c.numpy(), 2*i)
assert_jit_cache_len(add2, 1)
def test_jit_shape_mismatch(self):
@TinyJit
def add(a, b): return (a+b).realize()
for _ in range(5):
a = Tensor.randn(10, 10)
b = Tensor.randn(10, 10)
add(a, b)
bad = Tensor.randn(20, 20)
with self.assertRaises(JitError):
add(a, bad)
def test_jit_shape_views_mismatch(self):
@TinyJit
def add(a): return (a+1).realize()
with self.assertRaises(JitError):
for i in range(1,5):
# a has an offset that the kernel doesn't know about
a = Tensor.randn(10, 10).realize()[:, i:i+2]
add(a)
def test_jit_duplicate_fail(self):
# the jit doesn't support duplicate arguments
@TinyJit
def add(a, b): return (a+b).realize()
a = Tensor.randn(10, 10)
with self.assertRaises(JitError):
add(a, a)
def test_kwargs_jit(self):
@TinyJit
def add_kwargs(first, second): return (first+second).realize()
for _ in range(5):
a = Tensor.randn(10, 10)
b = Tensor.randn(10, 10)
c = add_kwargs(first=a, second=b)
np.testing.assert_allclose(c.numpy(), a.numpy()+b.numpy(), atol=1e-4, rtol=1e-5)
assert_jit_cache_len(add_kwargs, 1)
def test_array_jit(self):
@TinyJit
def add_array(a, arr): return (a+arr[0]).realize()
for _ in range(5):
a, b = Tensor.randn(10, 10).realize(), Tensor.randn(10, 10).realize()
np.testing.assert_allclose(add_array(a, [b]).numpy(), a.numpy()+b.numpy(), atol=1e-4, rtol=1e-5)
assert_jit_cache_len(add_array, 1)
def test_method_jit(self):
class Fun:
def __init__(self):
self.a = Tensor.randn(10, 10)
@TinyJit
def __call__(self, b:Tensor) -> Tensor:
return (self.a+b).realize()
fun = Fun()
for _ in range(5):
b = Tensor.randn(10, 10)
c = fun(b)
np.testing.assert_allclose(c.numpy(), fun.a.numpy()+b.numpy(), atol=1e-4, rtol=1e-5)
assert_jit_cache_len(fun.__call__.func.__self__, 1)
def test_jit_size1_input(self):
@TinyJit
def f(a, b): return (a+b).realize()
a = Tensor([1, 2, 3])
for i in range(5):
np.testing.assert_allclose(f(a, Tensor([i])).numpy(), (a+i).numpy(), atol=1e-4, rtol=1e-5)
assert_jit_cache_len(f, 1)
def test_jit_output_non_tensor_fail(self):
@TinyJit
def f(a, b, i): return (a+b).realize(), i
with self.assertRaises(JitError):
for i in range(3):
f(Tensor.randn(10, 10), Tensor.randn(10, 10), i)
def test_jit_random_regen(self):
def f(a, b):
rn = Tensor.randn(*a.shape)
return ((a+b)*rn).realize()
a = Tensor.randn(10, 10).realize() # realize these before resetting the random seed
b = Tensor.randn(10, 10).realize()
Tensor.manual_seed(1234)
jf = TinyJit(f)
res = set()
for _ in range(5):
o1 = jf(a, b)
res.add(o1.numpy()[0][0])
assert len(res) == 5, "All values should be different, rand works in jit."
Tensor.manual_seed(1234)
jf2 = TinyJit(f)
res2 = set()
for _ in range(5):
o1 = jf2(a, b)
res2.add(o1.numpy()[0][0])
assert len(res2) == 5, "All values should be different, rand works in jit."
assert res == res2, "Jit rand is not reproducible with the same seed"
Tensor.manual_seed(3421)
jf3 = TinyJit(f)
res3 = set()
for _ in range(5):
o1 = jf3(a, b)
res3.add(o1.numpy()[0][0])
assert len(res3) == 5, "All values should be different, rand works in jit."
assert res3 != res2, "Jit rand is diff with diff seeds"
def test_jit_v_nojit_random_regen(self):
def f(a, b):
rn = Tensor.randn(*a.shape)
rn = rn * a
rn2 = Tensor.randn(*a.shape)
rn2 = rn2 * b
rn = rn + rn2
rn2 = rn2 + Tensor.randn(*a.shape)
return ((a+b)*rn).realize(), ((a+b)*rn2).realize()
Tensor.manual_seed(0)
a = Tensor.randn(10, 10).realize() # realize these before resetting the random seed
b = Tensor.randn(10, 10).realize()
Tensor.manual_seed(1234)
without_jit = set()
for _ in range(5):
o1, o2 = f(a, b)
without_jit.add(o1.numpy()[0][0])
without_jit.add(o2.numpy()[0][0])
assert len(without_jit) == 10, "All values should be different."
Tensor.manual_seed(1234)
jf = TinyJit(f)
with_jit = set()
for _ in range(5):
o1, o2 = jf(a, b)
with_jit.add(o1.numpy()[0][0])
with_jit.add(o2.numpy()[0][0])
assert len(with_jit) == 10, "All values should be different."
assert with_jit == without_jit, "jit and non-jit should produce the same random values with the same seed"
def test_jit_multiple_random_regen(self):
def f(a, b):
rn = Tensor.randn(*a.shape)
rn = rn * a
rn2 = Tensor.randn(*a.shape)
rn2 = rn2 * b
rn = rn + rn2
rn2 = rn2 + Tensor.randn(*a.shape)
return ((a+b)*rn).realize(), ((a+b)*rn2).realize()
a = Tensor.randn(10, 10).realize() # realize these before resetting the random seed
b = Tensor.randn(10, 10).realize()
Tensor.manual_seed(1234)
jf = TinyJit(f)
res = set()
for _ in range(5):
o1, o2 = jf(a, b)
res.add(o1.numpy()[0][0])
res.add(o2.numpy()[0][0])
assert len(res) == 10, "All values should be different, rand works in jit."
Tensor.manual_seed(1234)
jf2 = TinyJit(f)
res2 = set()
for _ in range(5):
o1, o2 = jf2(a, b)
res2.add(o1.numpy()[0][0])
res2.add(o2.numpy()[0][0])
assert len(res2) == 10, "All values should be different, rand works in jit."
assert res == res2, "Jit rand is not reproducible with the same seed"
Tensor.manual_seed(3421)
jf3 = TinyJit(f)
res3 = set()
for _ in range(5):
o1, o2 = jf3(a, b)
res3.add(o1.numpy()[0][0])
res3.add(o2.numpy()[0][0])
assert len(res3) == 10, "All values should be different, rand works in jit."
assert res3 != res2, "Jit rand is diff with diff seeds"
def test_jit_random_after_unrealized_random(self):
@TinyJit
def f(): return Tensor.rand()
Tensor.manual_seed(1234)
Tensor.rand()
res = [f().numpy() for _ in range(3)]
assert res[1] != res[2]
def test_jit_realization_and_sampling(self):
w = Tensor.eye(5)
@TinyJit
def foo (x): return w.dot(x).realize()
arg = [
Tensor([1,2,3,4,5]),
Tensor([1,3,3,4,6]),
Tensor([1,2,5,4,7]),
Tensor([0,2,3,1,0]),
]
Y = [foo(e).numpy() for e in arg]
foo(Tensor([7,7,7,7,7]))
want = [[1., 2., 3., 4., 5.],
[1., 3., 3., 4., 6.],
[1., 2., 5., 4., 7.],
[0., 2., 3., 1., 0.]]
np.testing.assert_allclose(want, Y)
def test_jit_buffer_behavior(self):
@TinyJit
def foo(x) -> Tensor: return x.sum().realize()
result_1 = foo(Tensor([1] * 2))
result_2 = foo(Tensor([2] * 2))
result_3 = foo(Tensor([3] * 2))
# expect the buffer to share underlying buffer
np.testing.assert_allclose(result_1.numpy(), [2], atol=1e-4, rtol=1e-5)
np.testing.assert_allclose(result_2.numpy(), [6], atol=1e-4, rtol=1e-5)
np.testing.assert_allclose(result_3.numpy(), [6], atol=1e-4, rtol=1e-5)
def test_jit_output_clone(self):
@TinyJit
def f(x:Tensor) -> Tensor: return (x + 1).realize()
f(Tensor([0.0]))
f(Tensor([0.0]))
a = f(Tensor([1.0])).clone().realize()
b = f(Tensor([2.0]))
assert abs((a - b).item()) > 0.5
def test_jit_init_empty(self):
@TinyJit
def f(x:Tensor) -> Tensor: return (x + 1).realize()
f(Tensor.empty(1))
f(Tensor.empty(1))
# scalar const input is not allowed
with self.assertRaises(JitError):
f(Tensor(2.0)).item()
# self.assertEqual(f(Tensor([2.0])).item(), 1.0) # TODO: wrong output, should be 3.0. currently depends on empty value
def test_jit_const_input(self):
@TinyJit
def f(x:Tensor) -> Tensor: return (x + 1).realize()
with self.assertRaises(JitError):
f(Tensor(UOp.const(2.0).cast(dtypes.float))).item()
def test_jit_deviceless_compute_input(self):
@TinyJit
def f(x:Tensor) -> Tensor: return (x + 1).realize()
with self.assertRaises(JitError):
f(Tensor(UOp.const(2.0).cast(dtypes.float) + UOp.const(1.0).cast(dtypes.float))).item()
def test_jit_init_empty_alt(self):
@TinyJit
def f(a:Tensor, b:Tensor) -> Tensor: return b.assign(a+1)
for i in range(4):
a = Tensor([i])
b = Tensor.empty_like(a)
c = f(a, b)
self.assertEqual(c.item(), i+1)
def test_jit_lazy_grad_after_replay(self):
# the lazy .grad created during capture is read outside the JIT, the memory planner must not suballocate its buffers (issue #16571)
from tinygrad import nn
def step(conv, x, y):
out = conv(x.permute(0, 3, 1, 2).contiguous()).relu().flatten(1)
loss = (out * y).sum(axis=1) # per-example loss
loss.sum().backward()
conv.weight.grad = None
(loss * 0.5).sum().backward()
return loss.mean().realize()
Tensor.manual_seed(42)
conv = nn.Conv2d(3, 4, kernel_size=3, padding=1)
x, y = Tensor.randn(4, 8, 8, 3).realize(), Tensor.randn(4, 4*8*8).realize()
step(conv, x, y)
ref = conv.weight.grad.numpy()
jit_step = TinyJit(step)
for _ in range(4):
jit_step(conv, x, y)
np.testing.assert_allclose(conv.weight.grad.numpy(), ref, atol=1e-4, rtol=1e-5)
class TestJitPrune(unittest.TestCase):
def test_simple_prune(self):
weights = Tensor.rand(16).realize()
def w2(x) -> Tensor: return (weights*2).contiguous() + x
w2_noprune = TinyJit(w2)
w2_prune = TinyJit(w2, prune=True)
for _ in range(3):
a = Tensor.rand(16).realize()
out = w2_noprune(a)
np.testing.assert_allclose(out.tolist(), [x*2+y for x,y in zip(weights.tolist(), a.tolist())])
assert_jit_cache_len(w2_noprune, 2)
for _ in range(3):
a = Tensor.rand(16).realize()
out = w2_prune(a)
np.testing.assert_allclose(out.tolist(), [x*2+y for x,y in zip(weights.tolist(), a.tolist())])
assert_jit_cache_len(w2_prune, 1)
class TestJitInsideJit(unittest.TestCase):
def test_jit_jit_error(self):
@TinyJit
def f(t): return t + 1
@TinyJit
def g(t): return f(t) * 3
# NOTE: first does not raise
g(Tensor([1])).realize()
with self.assertRaisesRegex(RuntimeError, "having TinyJit inside another TinyJit is not supported"):
g(Tensor([1])).realize()
class TestJitRandom(unittest.TestCase):
def test_jit_rangeify(self):
tst = {0:[], 1:[]}
for r in [0,1]:
Tensor.manual_seed(1337)
with Context(JIT=r):
_ = Tensor.randint(4, high=3)
# this second one makes the behavior different
_ = Tensor.randint(4, high=3)
@TinyJit
def f(): return Tensor.randint(20, high=5)
for _ in range(5): tst[r].append(f().tolist())
for i, (t0, t1) in enumerate(zip(tst[0], tst[1])):
self.assertListEqual(t0, t1, msg=f"mismatch at list {i}")
if __name__ == '__main__':
unittest.main()

78
artifacts/package_sources/tinygrad/test/unit/test_jit_cases.py Normal file
View File

@@ -0,0 +1,78 @@
import unittest
from tinygrad import TinyJit, Tensor
# The JIT functions as a "capturing" JIT.
# Whatever kernels ran in the JIT the second run through the function will be the kernels that will run from then on.
# Explicit inputs to the function are updated in the JIT graph to the new inputs.
# JITs have four tensor types
# 1. Tensors that are explicit in the input, aka what's passed in. TODO: support lists/dicts/classes, anything get_state works on
# 2. Tensors that are explicit in the output, aka what's returned. TODO: same as above
# 3. Tensors that are implicit in the input as a closure.
# 4. Tensors that are implicit in the output because they were assigned to and realized.
# explicit inputs and outputs are realized on their way in and out of the JIT
# there's a whole bunch of edge cases and weirdness here that needs to be tested and clarified.
class TestJitCases(unittest.TestCase):
def test_explicit(self):
# this function has an explicit input and an explicit output
@TinyJit
def f(x:Tensor):
ret:Tensor = x*2
return ret
for i in range(5):
out = f(Tensor([i]))
self.assertEqual(out.item(), i*2)
def test_implicit_input(self):
# x is the implicit input (like a weight)
x = Tensor([0])
# this function has an implicit input and an explicit output
@TinyJit
def f():
ret:Tensor = x*2
return ret
for i in range(5):
# NOTE: this must be realized here, otherwise the update doesn't happen
# if we were explicitly tracking the implicit input Tensors, we might not need this realize
x.assign(Tensor([i])).realize()
out = f()
self.assertEqual(out.item(), i*2)
def test_implicit_output(self):
# out is the implicit output (it's assigned to)
out = Tensor([0])
# this function has an explicit input and an implicit output
@TinyJit
def f(x:Tensor):
# NOTE: this must be realized here
# if we were explicitly tracking the implicit output Tensors, we might not need this realize
out.assign(x*2).realize()
for i in range(5):
f(Tensor([i]))
self.assertEqual(out.item(), i*2)
def test_implicit_io(self):
# x is the implicit input (like a weight)
# out is the implicit output (it's assigned to)
x = Tensor([0])
out = Tensor([0])
# this function has an implicit input and an implicit output
@TinyJit
def f():
out.assign(x*2).realize() # NOTE: this must be realized here
for i in range(5):
x.assign(Tensor([i])).realize()
f()
self.assertEqual(out.item(), i*2)
if __name__ == '__main__':
unittest.main()

414
artifacts/package_sources/tinygrad/test/unit/test_jit_footguns.py Normal file
View File

@@ -0,0 +1,414 @@
#!/usr/bin/env python
"""
JIT Footguns: Documenting unexpected behavior changes when using @TinyJit
Each test shows behavior that works without JIT but changes with JIT.
Comments marked "should be X!" indicate the intuitively expected value.
SILENT MISMATCHES (highest priority - wrong results, no error):
class_method_shared_across_instances EASY could check if first arg is self and warn
slice_assign_requires_realize MED assign graph not connected to read during JIT replay
output_buffer_reuse MED performance tradeoff, could add option or better docs
symbolic_pad_view_frozen MED pad view BIND values baked in at capture time
python_constants_frozen HARD inherent to tracing JITs
conditional_branches_frozen HARD inherent to tracing JITs
ERRORS RAISED (lower priority - at least users know):
item_bakes_in_values EASY raises JitError if .item()/.data() accessed during capture
unrealized_const_input_error EASY raises JitError for unrealized const inputs
non_tensor_outputs_error EASY raises JitError if return contains non-Tensor values
positional_kwargs_cannot_mix EASY normalize positional args to kwargs using function signature
duplicate_inputs_fail MED would need to handle aliasing in input_replace
nested_jit_fails_on_second_call MED could fail on first call instead of second
"""
import unittest
import numpy as np
from tinygrad import Tensor, TinyJit, Device
from tinygrad.engine.jit import JitError
from tinygrad.helpers import JIT
class TestJitFootguns(unittest.TestCase):
def test_output_buffer_reuse(self):
"""Output tensors share buffer after capture - old references get overwritten."""
@TinyJit
def f(x): return x.sum().realize()
r1 = f(Tensor([1, 1])) # warmup
r2 = f(Tensor([2, 2])) # capture
r3 = f(Tensor([3, 3])) # jit exec
self.assertEqual(r1.item(), 2) # warmup result independent
self.assertEqual(r3.item(), 6) # latest is correct
self.assertEqual(r2.item(), 6) # should be 4! (overwritten by r3)
def test_output_buffer_workaround(self):
"""Use .clone().realize() to get independent copies."""
@TinyJit
def f(x): return x.sum().realize()
r1 = f(Tensor([1, 1])).clone().realize()
r2 = f(Tensor([2, 2])).clone().realize()
r3 = f(Tensor([3, 3])).clone().realize()
self.assertEqual([r1.item(), r2.item(), r3.item()], [2, 4, 6])
def test_graph_input_output_aliasing(self):
"""Test that JIT handles input=output aliasing correctly, simulating LLM generate pattern.
The LLM generate pattern:
1. First "session": multiple iterations where output becomes next input
2. Second "session": starts with a NEW input tensor (not the previous output)
The bug: GraphRunner computes input_replace during _first_run. If at that time input buffer == output buffer
(aliasing), it incorrectly includes the output position in input_replace. Later, when a DIFFERENT input
is passed, the output position gets overwritten with the input, corrupting the computation.
This requires multiple kernels to trigger because single-kernel JITs don't get graphed ("only one kernel doesn't graph").
"""
if Device[Device.DEFAULT].graph is None or JIT != 1:
self.skipTest("test requires JIT graph support")
# Multiple operations to create multiple kernels that get batched into a GraphRunner
@TinyJit
def step(x):
y = (x + 1).realize() # kernel 1
z = (y * 2).realize() # kernel 2
return z
# Phase 1: warmup and capture
a = Tensor([10]).contiguous().realize()
step(a) # warmup (cnt=0)
b = Tensor([20]).contiguous().realize()
x = step(b) # capture (cnt=1), x = (20+1)*2 = 42
# Phase 2: first "session" - iterations where output becomes input (triggers _first_run with aliasing)
for _ in range(3):
x = step(x) # (42+1)*2=86, (86+1)*2=174, (174+1)*2=350
self.assertEqual(x.item(), 350)
# Phase 3: second "session" - NEW input tensor (simulates new generate() call)
# The bug: GraphRunner's input_replace incorrectly includes the output position
# When new input y is passed, it overwrites the output buffer, using old value (350) instead of new (100)
y = Tensor([100]).contiguous().realize()
for _ in range(3):
y = step(y) # should be (100+1)*2=202, (202+1)*2=406, (406+1)*2=814
self.assertEqual(y.item(), 814) # fails with 1406 if bug exists (uses 350 instead of 100)
def test_multiple_outputs_same_intermediate(self):
"""Multiple outputs derived from the same intermediate - JIT copies aliased inputs to prevent hazard."""
@TinyJit
def f(buf, frame):
new_buf = buf[1:].cat(frame, dim=0)
return new_buf.contiguous(), new_buf[:1].contiguous()
buf = Tensor([[0], [1], [2]]).contiguous().realize()
for i in range(4):
frame = Tensor([[10+i]]).contiguous().realize()
expected_first = buf[1:2].numpy().item()
new_buf, first = f(buf, frame)
self.assertEqual(first.numpy().item(), expected_first)
buf = new_buf
def test_intra_kernel_output_input_aliasing(self):
"""JIT must copy aliased input when output buffer is fed back as input (read-write race in same kernel)."""
N = 1 << 20
f = TinyJit(lambda buf, new: buf[N//2:].cat(new), prune=True)
buf = Tensor.zeros(N, dtype='int32').contiguous().realize()
for i in range(10):
buf = f(buf, Tensor(np.ones(N//2, dtype=np.int32)*(i+1)))
np.testing.assert_array_equal(buf[:N//2].numpy(), np.full(N//2, i, dtype=np.int32))
def test_slice_assign_works_without_realize(self):
"""Slice assign then read from same buffer - pending assigns are side-realized."""
from tinygrad import Variable
v_pos = Variable("pos", 0, 3)
cache = Tensor.zeros(4, 4).contiguous().realize()
@TinyJit
def f(pos):
cache[pos:pos+1, :].assign(Tensor.ones(1, 4))
return cache.sum().realize()
for i in range(4):
cache.assign(Tensor.zeros(4, 4)).realize()
self.assertEqual(f(v_pos.bind(i)).item(), 4.0)
def test_symbolic_pad_view_frozen(self):
"""Symbolic pad view has BIND values baked in at capture time. TODO: pad should be captured in jit."""
from tinygrad import Variable
a = Tensor.rand(3, 10).realize()
# broken: pad is a view, BIND values frozen at capture (i=2)
@TinyJit
def f_broken(a): return (a+1).pad((None, (0, 10-a.shape[1]))).realize()
for i in range(1, 5): f_broken(a[:, :Variable("i", 1, 10).bind(i)])
self.assertEqual(int((f_broken(a[:, :Variable("i", 1, 10).bind(4)])[0] != 0).sum().item()), 2) # should be 4!
# workaround: contiguous fuses pad into kernel
@TinyJit
def f_fixed(a): return (a+1).pad((None, (0, 10-a.shape[1]))).contiguous().realize()
for i in range(1, 5): f_fixed(a[:, :Variable("i", 1, 10).bind(i)])
self.assertEqual(int((f_fixed(a[:, :Variable("i", 1, 10).bind(4)])[0] != 0).sum().item()), 4)
def test_non_tensor_outputs_error(self):
@TinyJit
def f(x, mult): return (x * 2).realize(), mult * 10
with self.assertRaises(JitError):
for i in range(3): f(Tensor([i]), i)
def test_duplicate_inputs_fail(self):
"""JIT cannot handle the same tensor passed as multiple arguments."""
@TinyJit
def f(a, b): return (a + b).realize()
x = Tensor([1, 2, 3])
with self.assertRaises(JitError):
f(x, x)
def test_tensors_in_containers(self):
@TinyJit
def f(a, arr): return (a + arr[0]).realize()
for i in range(4):
a, b = Tensor([1, 1, 1]).realize(), Tensor([i, i, i]).realize()
np.testing.assert_array_equal(f(a, [b]).numpy(), [1+i, 1+i, 1+i])
def test_nested_jit_fails_on_second_call(self):
"""Nested JIT works on first call but fails on second."""
@TinyJit
def inner(t): return t + 1
@TinyJit
def outer(t): return inner(t) * 3
self.assertEqual(outer(Tensor([1])).realize().item(), 6) # works!
with self.assertRaises(RuntimeError):
outer(Tensor([2])).realize() # fails
def test_implicit_inputs_need_realize(self):
"""Closure tensors must be realized before JIT call."""
x = Tensor([0])
@TinyJit
def f(): return (x * 2).realize()
for i in range(5):
x.assign(Tensor([i])).realize() # must realize!
self.assertEqual(f().item(), i * 2)
def test_views_with_different_offsets_fail(self):
"""JIT requires consistent tensor views across calls."""
@TinyJit
def f(a): return (a + 1).realize()
base = Tensor.randn(10, 10).realize()
with self.assertRaises(JitError):
for i in range(1, 5):
f(base[:, i:i+2]) # different offset each time
def test_shape_change_after_capture_fails(self):
"""Shapes are locked at capture time."""
@TinyJit
def f(a, b): return (a + b).realize()
f(Tensor.randn(10, 10), Tensor.randn(10, 10)) # warmup
f(Tensor.randn(10, 10), Tensor.randn(10, 10)) # capture
with self.assertRaises(JitError):
f(Tensor.randn(20, 20), Tensor.randn(20, 20))
def test_python_constants_frozen(self):
"""Python variables inside JIT use capture-time values."""
mult = 1
@TinyJit
def f(x): return (x * mult).realize()
results = []
for i in range(5):
mult = i + 1
results.append(f(Tensor([10])).item())
self.assertEqual(results[0], 10) # warmup, mult=1
self.assertEqual(results[1], 20) # capture, mult=2
self.assertEqual(results[2], 20) # should be 30!
self.assertEqual(results[3], 20) # should be 40!
def test_unrealized_const_input_error(self):
"""Const tensors have no buffer to replace, so JIT raises an error. Even explicit .realize() doesn't help."""
@TinyJit
def f(a, b): return (a * b).realize()
# unrealized const fails
with self.assertRaises(JitError):
f(Tensor([1, 2, 3]).realize(), Tensor(2))
# explicit .realize() on const still fails - const cannot be realized to have a buffer
@TinyJit
def g(a, b): return (a * b).realize()
with self.assertRaises(JitError):
g(Tensor([1, 2, 3]).realize(), Tensor(2).realize())
def test_conditional_branches_frozen(self):
"""Only the branch taken during capture runs thereafter."""
@TinyJit
def f(x, use_square):
if use_square:
return (x * x).realize()
return (x * 2).realize()
f(Tensor([3]), True) # warmup
f(Tensor([3]), False) # capture (False branch)
result = f(Tensor([3]), True) # passing True but False branch runs
self.assertEqual(result.item(), 6) # should be 9!
def test_positional_kwargs_cannot_mix(self):
"""Must use same calling convention after capture."""
@TinyJit
def f(a, b): return (a + b).realize()
f(Tensor([1]), Tensor([2])) # warmup with positional
f(Tensor([1]), Tensor([2])) # capture with positional
with self.assertRaises(JitError):
f(a=Tensor([3]), b=Tensor([4])) # kwargs fail
def test_class_method_shared_across_instances(self):
"""JIT on instance methods is shared at class level."""
class Model:
def __init__(self, scale):
self.scale = Tensor([scale])
@TinyJit
def forward(self, x):
return (x * self.scale).realize()
m1, m2 = Model(2), Model(3)
m1.forward(Tensor([5])) # warmup
m1.forward(Tensor([5])) # capture with m1.scale=2
self.assertEqual(m1.forward(Tensor([5])).item(), 10)
self.assertEqual(m2.forward(Tensor([5])).item(), 10) # should be 15!
def test_side_effects_only_during_capture(self):
"""Function body not executed during JIT replay."""
call_count = [0]
@TinyJit
def f(x):
call_count[0] += 1
return (x * 2).realize()
f(Tensor([1])) # warmup
f(Tensor([2])) # capture
self.assertEqual(call_count[0], 2)
f(Tensor([3]))
f(Tensor([4]))
f(Tensor([5]))
self.assertEqual(call_count[0], 2) # still 2, not 5!
def test_nothing_realized_fails(self):
"""Must JIT at least one kernel."""
@TinyJit
def f(a, b): return None
with self.assertRaises(JitError):
for _ in range(3):
f(Tensor([1]), Tensor([2]))
def test_item_creates_unrealized_return(self):
""".item() in shape computation raises error during JIT capture."""
@TinyJit
def f(x): return Tensor.zeros(x.sum().item())
f(Tensor([1, 1, 1])) # warmup
with self.assertRaises(JitError):
f(Tensor([1, 1, 1])) # capture - .item() raises
def test_item_bakes_in_values(self):
""".item() during JIT capture raises error (would bake in value)."""
@TinyJit
def f(x, mask): return x.masked_select(mask)
f(Tensor([1, 2, 3, 4]), Tensor([True, False, True, False])) # warmup
with self.assertRaises(JitError):
f(Tensor([1, 2, 3, 4]), Tensor([True, False, True, False])) # capture - .item() raises
def test_masked_select_static_size_jittable(self):
@TinyJit
def f(x, mask): return x.masked_select(mask, size=4, fill_value=-1).realize()
for _ in range(3):
np.testing.assert_equal(f(Tensor([1, 2, 3, 4]), Tensor([True, False, True, False])).numpy(), [1, 3, -1, -1])
np.testing.assert_equal(f(Tensor([5, 6, 7, 8]), Tensor([False, True, True, True])).numpy(), [6, 7, 8, -1])
np.testing.assert_equal(f(Tensor([9, 8, 7, 6]), Tensor([True, True, True, True])).numpy(), [9, 8, 7, 6])
np.testing.assert_equal(f(Tensor([1, 1, 1, 1]), Tensor([False, False, False, False])).numpy(), [-1, -1, -1, -1])
def test_nonzero_static_size_jittable(self):
@TinyJit
def f(x): return x.nonzero(size=3, fill_value=-1).realize()
for _ in range(3):
np.testing.assert_equal(f(Tensor([1, 0, 2, 0, 3])).numpy(), [[0], [2], [4]])
np.testing.assert_equal(f(Tensor([0, 0, 5, 0, 0])).numpy(), [[2], [-1], [-1]])
np.testing.assert_equal(f(Tensor([0, 0, 0, 0, 0])).numpy(), [[-1], [-1], [-1]])
def test_tolist_bakes_in_values(self):
""".tolist() raises error during JIT capture (would bake in values)."""
@TinyJit
def f(x): return Tensor(x.tolist())
f(Tensor([1, 2, 3])) # warmup
with self.assertRaises(JitError):
f(Tensor([1, 2, 3])) # capture - .tolist() raises
class TestJitCorrectBehavior(unittest.TestCase):
"""Behaviors that work correctly - documented for clarity."""
def test_random_regenerates(self):
"""Random tensors regenerate each call."""
@TinyJit
def f(x):
return (x + Tensor.rand(3)).realize()
f(Tensor([0, 0, 0])) # warmup
f(Tensor([0, 0, 0])) # capture
results = {tuple(f(Tensor([0, 0, 0])).numpy().tolist()) for _ in range(5)}
self.assertEqual(len(results), 5)
def test_unrealized_return_auto_realized(self):
"""Unrealized return tensors are auto-realized."""
@TinyJit
def f(a, b): return a + b # no explicit realize
for _ in range(5):
a, b = Tensor.randn(10), Tensor.randn(10)
np.testing.assert_allclose(f(a, b).numpy(), a.numpy() + b.numpy(), atol=1e-5)
def test_kwargs_order_doesnt_matter(self):
"""Kwargs are sorted by name, so order doesn't matter."""
@TinyJit
def f(first, second): return (first / second).realize()
for _ in range(3):
a, b = Tensor.randn(10), Tensor.randn(10) + 1
np.testing.assert_allclose(f(second=b, first=a).numpy(), a.numpy() / b.numpy(), atol=1e-4)
np.testing.assert_allclose(f(first=a, second=b).numpy(), a.numpy() / b.numpy(), atol=1e-4)
def test_input_mutation_consistent(self):
"""Input mutation via assign works consistently."""
@TinyJit
def f(x):
x += 1
x.realize()
return x
a = Tensor([0]).contiguous().realize()
for _ in range(5):
f(a)
self.assertEqual(a.item(), 5)
if __name__ == '__main__':
unittest.main()

116
artifacts/package_sources/tinygrad/test/unit/test_linalg.py Normal file
View File

@@ -0,0 +1,116 @@
import unittest, functools
from tinygrad import Tensor, Context
import numpy as np
def orthogonality_helper(A:Tensor, tolerance=1e-5):
b_shape,m = A.shape[0:-2],A.shape[-2] #outer dimension should be the dim along orthogonality
A_identity = (Tensor.eye(m).reshape((1,)*len(b_shape)+(m,m)).expand(b_shape+(m,m)))
np.testing.assert_allclose((A @ A.transpose(-2,-1)).numpy(),A_identity.numpy(),atol=tolerance,rtol=tolerance)
def reconstruction_helper(A:list[Tensor],B:Tensor, tolerance=1e-5):
reconstructed_tensor = functools.reduce(Tensor.matmul, A)
np.testing.assert_allclose(reconstructed_tensor.numpy(),B.numpy(),atol=tolerance,rtol=tolerance)
class TestLinAlg(unittest.TestCase):
def test_svd_general(self):
sizes = [(2,2),(5,3),(3,5),(3,4,4),(2,2,2,2,3)]
for size in sizes:
a = Tensor.randn(size).realize()
U,S,V = a.svd()
Tensor.realize(U,S,V)
b_shape,m,n = size[0:-2],size[-2],size[-1]
k = min(m,n)
s_diag = (S.unsqueeze(-2) * Tensor.eye(k).reshape((1,) * len(b_shape) + (k,k)))
s_diag = s_diag.expand(b_shape + (k,k)).pad(tuple([None]*len(b_shape) + [(0,m-k), (0,n-k)]))
orthogonality_helper(U)
orthogonality_helper(V)
reconstruction_helper([U,s_diag,V],a)
def _test_svd_nonfull(self, size):
with Context(CHECK_OOB=0): # sometimes this is slow in CI
a = Tensor.randn(size).realize()
U,S,V = a.svd(full_matrices=False)
Tensor.realize(U,S,V)
b_shape,m,n = size[0:-2],size[-2],size[-1]
k = min(m,n)
s_diag = (S.unsqueeze(-2) * Tensor.eye(k).reshape((1,) * len(b_shape) + (k,k)).expand(b_shape + (k,k)))
#reduced U,V is only orthogonal along smaller dim
if (m < n): orthogonality_helper(U),orthogonality_helper(V)
else: orthogonality_helper(U.transpose(-2,-1)),orthogonality_helper(V.transpose(-2,-1))
reconstruction_helper([U,s_diag,V],a)
# faster for parallel pytest
def test_svd_nonfull_2_2(self): self._test_svd_nonfull((2,2))
def test_svd_nonfull_5_3(self): self._test_svd_nonfull((5,3))
def test_svd_nonfull_3_5(self): self._test_svd_nonfull((3,5))
def test_svd_nonfull_2_2_2_2_3(self): self._test_svd_nonfull((2,2,2,2,3))
@unittest.skip("very big. recommend wrapping with TinyJit around inner function")
def test_svd_large(self):
size = (1024,1024)
a = Tensor.randn(size).realize()
U,S,V = a.svd()
b_shape,m,n = size[0:-2],size[-2],size[-1]
k = min(m,n)
s_diag = (S.unsqueeze(-2) * Tensor.eye(k).reshape((1,) * len(b_shape) + (k,k)))
s_diag = s_diag.expand(b_shape + (k,k)).pad(tuple([None]*len(b_shape) + [(0,m-k), (0,n-k)]))
orthogonality_helper(U,tolerance=1e-3)
orthogonality_helper(V,tolerance=1e-3)
reconstruction_helper([U,s_diag,V],a,tolerance=1e-3)
def test_qr_general(self):
sizes = [(3,3),(3,6),(6,3),(2,2,2,2,2)]
for size in sizes:
a = Tensor.randn(size).realize()
Q,R = a.qr()
Tensor.realize(Q,R)
orthogonality_helper(Q)
reconstruction_helper([Q,R],a)
def test_qr_zero_column(self):
a = Tensor([[0.0, 1.0], [0.0, 2.0]]).realize()
Q,R = a.qr()
assert not np.isnan(Q.numpy()).any()
assert not np.isnan(R.numpy()).any()
orthogonality_helper(Q)
reconstruction_helper([Q,R], a)
def test_svd_identity(self):
for a in (Tensor.eye(2).clone(), Tensor.zeros(2, 2)):
a = a.realize()
U,S,V = a.svd()
Tensor.realize(U,S,V)
assert not np.isnan(U.numpy()).any()
assert not np.isnan(S.numpy()).any()
assert not np.isnan(V.numpy()).any()
s_diag = (S.unsqueeze(-2) * Tensor.eye(2))
reconstruction_helper([U, s_diag, V], a)
def test_svd_identity_4x4(self):
a = Tensor.eye(4).clone()
U,S,V = a.svd()
Tensor.realize(U,S,V)
assert not np.isnan(U.numpy()).any()
assert not np.isnan(S.numpy()).any()
assert not np.isnan(V.numpy()).any()
s_diag = (S.unsqueeze(-2) * Tensor.eye(4))
reconstruction_helper([U, s_diag, V], a)
def test_svd_rank1(self):
a = Tensor([[1.0, 1.0], [2.0, 2.0]]).realize()
U, S, V = a.svd()
np.testing.assert_allclose(S.numpy(), [np.sqrt(10), 0.0], atol=1e-4, rtol=1e-4)
reconstruction_helper([U, S.unsqueeze(-2) * Tensor.eye(2), V], a)
def test_newton_schulz(self):
coefficients = [(2, -1.5, 0.5), (2.0, -1.4, 0.2, 0.2)]#these params map to the sign function
sizes = [(2,2), (3,2), (2,3), (2,2,2)]
for coefs in coefficients:
for size in sizes:
a = Tensor.randn(size)
b = a.newton_schulz(steps=20, params=coefs, eps=0.0)
# ns(A) = U @ Vt -> (U @ Vt) @ (U @ Vt)t = I
orthogonality_helper(b if size[-1] > size[-2] else b.transpose(-2, -1), tolerance=1e-3)
if __name__ == "__main__":
unittest.main()

71
artifacts/package_sources/tinygrad/test/unit/test_llm_mla.py Normal file
View File

@@ -0,0 +1,71 @@
import unittest
import numpy as np
from tinygrad import Tensor
from tinygrad.llm.model import Transformer, TransformerConfig, apply_rope, MLATransformerBlock, precompute_freqs_cis
class TestMLA(unittest.TestCase):
def _make_config(self, **kwargs):
return TransformerConfig(**{
"num_blocks": 1, "dim": 64, "hidden_dim": 128, "n_heads": 4, "n_kv_heads": 1,
"norm_eps": 1e-5, "vocab_size": 100, "head_dim": 16, "rope_theta": 10000.0, "rope_dim": 8, "max_context": 32,
"kv_lora_rank": 16, "v_head_dim": 8,
} | kwargs)
def test_mla_attention_matches_naive(self):
config = self._make_config(max_context=16)
block = MLATransformerBlock(config)
c = config
B, T = 1, 4
q_nope_head_dim = c.head_dim - c.rope_dim
x = Tensor.randn(B, T, c.dim)
x_norm = block.attn_norm(x)
# --- Our absorbed implementation ---
q = block.attn_q(x_norm).reshape(B, T, c.n_heads, c.head_dim).transpose(1, 2)
q_nope, q_rope = q[..., :q_nope_head_dim], q[..., q_nope_head_dim:]
freqs = precompute_freqs_cis(c.rope_dim, 16, c.rope_theta)
q_rope = apply_rope(q_rope, freqs[0:T])
kv_a = block.attn_kv_a_mqa(x_norm)
c_kv = block.attn_kv_a_norm(kv_a[..., :c.kv_lora_rank])
k_rope = kv_a[..., c.kv_lora_rank:].reshape(B, T, 1, c.rope_dim).transpose(1, 2)
k_rope = apply_rope(k_rope, freqs[0:T])
# --- Naive (non-absorbed): expand K and V, do standard attention ---
k_nope_naive = c_kv.unsqueeze(1) @ block.attn_k_b["weight"] # (B, H, T, nope)
k_naive = k_nope_naive.cat(k_rope.expand(-1, c.n_heads, -1, -1), dim=-1) # (B, H, T, nope+rope)
v_naive = c_kv.unsqueeze(1) @ block.attn_v_b["weight"].transpose(-1, -2) # (B, H, T, v_dim)
q_naive = q_nope.cat(q_rope, dim=-1)
scale = 1.0 / c.head_dim ** 0.5
scores_naive = (q_naive @ k_naive.transpose(-1, -2)) * scale
# causal mask
mask = Tensor.full((1, 1, T, T), float("-inf")).triu(1)
attn_naive = (scores_naive + mask).softmax(-1) @ v_naive # (B, H, T, v_dim)
out_naive = block.attn_output(attn_naive.transpose(1, 2).reshape(B, T, -1))
# --- Absorbed: q_nope @ wk_b^T, then dot with compressed kv ---
q_nope_abs = q_nope @ block.attn_k_b["weight"].transpose(-1, -2) # (B, H, T, lora)
q_abs = q_nope_abs.cat(q_rope, dim=-1) # (B, H, T, lora+rope)
k_abs = c_kv.reshape(B, 1, T, c.kv_lora_rank).cat(k_rope.reshape(B, 1, T, c.rope_dim), dim=-1)
scores_abs = (q_abs @ k_abs.transpose(-1, -2)) * scale
attn_abs = (scores_abs + mask).softmax(-1)
# attn @ v_compressed @ wv_b
v_compressed = c_kv.reshape(B, 1, T, c.kv_lora_rank)
attn_abs_out = (attn_abs @ v_compressed) @ block.attn_v_b["weight"].transpose(-1, -2)
out_abs = block.attn_output(attn_abs_out.transpose(1, 2).reshape(B, T, -1))
# Compare
naive_np = out_naive.realize().numpy()
abs_np = out_abs.realize().numpy()
np.testing.assert_allclose(naive_np, abs_np, atol=1e-4, rtol=1e-4,
err_msg="Absorbed MLA should match naive MLA")
def test_shared_expert_gate_optional(self):
from tinygrad import nn
model = Transformer(self._make_config(num_experts=4, num_experts_per_tok=2, shared_expert_dim=32, shared_expert_gate=False))
self.assertNotIn('blk.0.ffn_gate_inp_shexp.weight', nn.state.get_state_dict(model))
out = model.blk[0]._feed_forward(Tensor.randn(1, 4, model.blk[0].config.dim))
self.assertEqual(out.shape, (1, 4, model.blk[0].config.dim))

100
artifacts/package_sources/tinygrad/test/unit/test_llm_moe.py Normal file
View File

@@ -0,0 +1,100 @@
import unittest
import numpy as np
from dataclasses import replace
from tinygrad import Tensor
from tinygrad.llm.model import TransformerBlock, TransformerConfig
def _moe_config(dim=8, hidden=16, n_heads=2, num_experts=4, num_experts_per_tok=2):
return TransformerConfig(
num_blocks=1, dim=dim, hidden_dim=hidden, n_heads=n_heads, n_kv_heads=n_heads,
norm_eps=1e-5, vocab_size=100, head_dim=dim//n_heads, rope_theta=10000,
rope_dim=dim//n_heads, v_head_dim=dim//n_heads, max_context=16,
num_experts=num_experts, num_experts_per_tok=num_experts_per_tok)
class TestMoEFeedForward(unittest.TestCase):
def test_moe_feed_forward(self):
dim, hidden, n_heads = 8, 16, 2
num_experts, k = 4, 2
block = TransformerBlock(_moe_config(dim, hidden, n_heads, num_experts, k))
# set up weights: gate scales by (expert_id+1), up/down are identity-ish, router picks experts 0,2
block.ffn_gate_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) * (i + 1) for i in range(num_experts)])
block.ffn_up_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) for _ in range(num_experts)])
block.ffn_down_exps.weight = Tensor.stack(*[Tensor.eye(dim, hidden) for _ in range(num_experts)])
block.ffn_gate_inp.weight = Tensor([[1, 0, 1, 0]] * dim).T # router strongly prefers experts 0 and 2
block.ffn_norm.weight = Tensor.ones(dim) # identity norm
# input of ones -> after norm still ~ones -> experts 0,2 selected -> weighted sum of silu outputs
h = Tensor.ones(1, 1, dim)
out = block._feed_forward(block.ffn_norm(h))
# expected moe_output ≈ avg(silu(1), silu(3))
expected = (Tensor([1.0]).silu().item() + Tensor([3.0]).silu().item()) / 2
np.testing.assert_allclose(out.numpy()[0, 0, 0], expected, rtol=1e-2)
def test_moe_feed_forward_batched(self):
dim, hidden, n_heads = 8, 16, 2
num_experts, k = 4, 2
block = TransformerBlock(_moe_config(dim, hidden, n_heads, num_experts, k))
# same setup as BS=1 test
block.ffn_gate_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) * (i + 1) for i in range(num_experts)])
block.ffn_up_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) for _ in range(num_experts)])
block.ffn_down_exps.weight = Tensor.stack(*[Tensor.eye(dim, hidden) for _ in range(num_experts)])
block.ffn_gate_inp.weight = Tensor([[1, 0, 1, 0]] * dim).T
block.ffn_norm.weight = Tensor.ones(dim)
# test with BS=2, T=3
h = Tensor.ones(2, 3, dim)
out = block._feed_forward(block.ffn_norm(h))
# all outputs should match the BS=1 expected value
expected = (Tensor([1.0]).silu().item() + Tensor([3.0]).silu().item()) / 2
np.testing.assert_allclose(out.numpy(), expected, rtol=1e-2)
def test_moe_feed_forward_norm_topk_prob(self):
dim, hidden, n_heads = 8, 16, 2
num_experts, k = 4, 2
block = TransformerBlock(replace(_moe_config(dim, hidden, n_heads, num_experts, k), norm_topk_prob=True))
block.ffn_gate_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) * (i + 1) for i in range(num_experts)])
block.ffn_up_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) for _ in range(num_experts)])
block.ffn_down_exps.weight = Tensor.stack(*[Tensor.eye(dim, hidden) for _ in range(num_experts)])
block.ffn_gate_inp.weight = Tensor([[0.1, 0, 0.1, 0]] * dim).T # equal top-2 experts, but only ~69% mass before renorm
block.ffn_norm.weight = Tensor.ones(dim)
h = Tensor.ones(1, 1, dim)
out = block._feed_forward(block.ffn_norm(h))
expected = (Tensor([1.0]).silu().item() + Tensor([3.0]).silu().item()) / 2
np.testing.assert_allclose(out.numpy()[0, 0, 0], expected, rtol=1e-2)
def test_moe_feed_forward_shared_expert(self):
dim, hidden, n_heads = 8, 16, 2
num_experts, k = 4, 2
block = TransformerBlock(replace(_moe_config(dim, hidden, n_heads, num_experts, k), shared_expert_dim=dim))
block.ffn_gate_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) * (i + 1) for i in range(num_experts)])
block.ffn_up_exps.weight = Tensor.stack(*[Tensor.eye(hidden, dim) for _ in range(num_experts)])
block.ffn_down_exps.weight = Tensor.stack(*[Tensor.eye(dim, hidden) for _ in range(num_experts)])
block.ffn_gate_inp.weight = Tensor([[1, 0, 1, 0]] * dim).T
block.ffn_gate_shexp.weight = Tensor.eye(dim) * 2
block.ffn_up_shexp.weight = Tensor.eye(dim)
block.ffn_down_shexp.weight = Tensor.eye(dim)
block.ffn_gate_inp_shexp["weight"] = Tensor.zeros(dim)
block.ffn_norm.weight = Tensor.ones(dim)
h = Tensor.ones(1, 1, dim)
out = block._feed_forward(block.ffn_norm(h))
moe_expected = (Tensor([1.0]).silu().item() + Tensor([3.0]).silu().item()) / 2
shared_expected = Tensor([2.0]).silu().item() * 0.5
expected = moe_expected + shared_expected
np.testing.assert_allclose(out.numpy(), expected, rtol=1e-2)
if __name__ == '__main__':
unittest.main()

194
artifacts/package_sources/tinygrad/test/unit/test_llm_server.py Normal file
View File

@@ -0,0 +1,194 @@
import unittest
from unittest.mock import patch
from tinygrad import Tensor, UOp
from tinygrad.schedule import schedule_cache
from tinygrad.llm.model import Transformer, TransformerConfig
from tinygrad.llm.serve import StreamRouter
TEST_CONFIG = TransformerConfig(num_blocks=1, dim=64, hidden_dim=128, n_heads=2, n_kv_heads=2,
norm_eps=1e-5, vocab_size=100, head_dim=32, rope_theta=10000.0, rope_dim=32, v_head_dim=32, max_context=32)
V_START_POS = UOp.variable("start_pos", 0, TEST_CONFIG.max_context-1)
V_TOKS = UOp.variable("toks", 1, 32) # 32 is the default chunk_size in generate
class TestTransformerGenerate(unittest.TestCase):
def test_warmup(self):
model, calls = Transformer(TEST_CONFIG), []
def generate(tokens):
calls.append(tokens)
yield from (1, 2)
with patch.object(model, "generate", generate): model.warmup()
self.assertEqual(calls, [[0], [0]])
def test_first_recurrent_generate_before_state_init(self):
model = Transformer(TEST_CONFIG)
model.has_recurrent_block = True
with patch.object(Transformer, '__call__', return_value=Tensor([[42]])):
self.assertEqual(next(model.generate([0])), 42)
def test_recurrent_live_state_reuse(self):
model = Transformer(TEST_CONFIG)
model.has_recurrent_block = True
model._cached_tokens = [1, 2, 3, 4, 5]
self.assertEqual(model.get_start_pos([1, 2, 3, 4, 5, 42, 10]), 5)
calls = []
def mock_call(self, tokens, start_pos, temperature, **kwargs):
calls.append((tokens.shape, start_pos))
return Tensor([[42]])
with patch.object(Transformer, '__call__', mock_call):
next(model.generate([1, 2, 3, 4, 5, 42, 10]))
self.assertEqual(calls, [((1, 1), V_START_POS.bind(5)), ((1, 1), V_START_POS.bind(6))])
def test_template_starts_reasoning(self):
router = StreamRouter(reasoning=True)
self.assertEqual(list(router.route("reasoning</think>answer")),
[("reasoning_content", "reasoning"), ("content", "answer")])
def test_kv_cache_reuse(self):
"""Test that generate reuses the KV cache when tokens extend the cached prefix."""
model = Transformer(TEST_CONFIG)
captured_inputs = []
def mock_call(self, tokens, start_pos, temperature, **kwargs):
captured_inputs.append((tokens.shape, start_pos))
return Tensor([[42]])
with patch.object(Transformer, '__call__', mock_call):
# first conversation: prefill 5 tokens + 1 decode
tokens = [1, 2, 3, 4, 5]
gen = model.generate(tokens)
next(gen) # prefill
next(gen) # decode
# second call extends the conversation — cached prefix should be reused
captured_inputs.clear()
tokens = [1, 2, 3, 4, 5, 42, 42, 10, 11, 12]
gen = model.generate(tokens)
next(gen)
# should process tokens[6:] = [42, 10, 11, 12] since first 6 have cached k/v
self.assertEqual(captured_inputs, [((1, V_TOKS.bind(4)), V_START_POS.bind(6))])
def test_kv_cache_invalidation(self):
"""Test that generate invalidates the KV cache when tokens diverge from the cached prefix."""
model = Transformer(TEST_CONFIG)
captured_inputs = []
def mock_call(self, tokens, start_pos, temperature, **kwargs):
captured_inputs.append((tokens.shape, start_pos))
return Tensor([[42]])
with patch.object(Transformer, '__call__', mock_call):
# first conversation
gen = model.generate([1, 2, 3, 4, 5])
next(gen)
# completely different prompt — KV cache should be invalidated
captured_inputs.clear()
gen = model.generate([10, 20, 30])
next(gen)
# should process all 3 tokens from start
self.assertEqual(captured_inputs, [((1, V_TOKS.bind(3)), V_START_POS.bind(0))])
def test_two_prompts_schedule_cache(self):
"""Third prompt should hit the schedule cache, not miss (first two warm up both jits: prefill + decode)."""
from dataclasses import replace
model = Transformer(replace(TEST_CONFIG, max_context=64))
# first two prompts warm up both jits (prefill + decode)
ids = list(range(1, 6))
gen = model.generate(ids)
for _ in range(3): next(gen)
ids += list(range(10, 15))
gen = model.generate(ids)
for _ in range(3): next(gen)
cache_size_after_warmup = len(schedule_cache)
# third prompt should reuse the same schedule cache entries, not create new ones
ids += list(range(20, 25))
gen = model.generate(ids)
for _ in range(3): next(gen)
self.assertEqual(cache_size_after_warmup, len(schedule_cache),
f"third prompt added {len(schedule_cache) - cache_size_after_warmup} new schedule cache entries (expected 0)")
def test_chunked_prefill(self):
"""When prompt > chunk_size, all chunks should be prefill"""
from tinygrad.uop.ops import resolve
from dataclasses import replace
model = Transformer(replace(TEST_CONFIG, max_context=64))
def get_prefill_flags(tokens, chunk_size):
is_prefill = []
def mock_call(self, tokens, start_pos, temperature, **kwargs):
is_prefill.append(resolve(tokens.shape[1] != 1))
return Tensor([[42]])
with patch.object(Transformer, '__call__', mock_call):
gen = model.generate(tokens, chunk_size=chunk_size)
for _ in range(3): next(gen)
model._cached_tokens = []
return is_prefill
# 8 tokens, chunk_size=4 -> 2 prefill chunks
self.assertEqual(get_prefill_flags(list(range(8)), 4), [True, True, False, False])
# 9 tokens, chunk_size=4 -> 3 prefill chunks (4+4+1)
self.assertEqual(get_prefill_flags(list(range(9)), 4), [True, True, True, False, False])
# 4 tokens, chunk_size=4 -> 1 prefill chunk
self.assertEqual(get_prefill_flags(list(range(4)), 4), [True, False, False])
def test_kv_cache_resume_matches_fresh(self):
model = Transformer(TEST_CONFIG)
# generate 2 tokens, then abandon
prompt = list(range(1, 6))
gen = model.generate(list(prompt))
out1, out2 = next(gen), next(gen)
# resume with conversation history + new user tokens appended
extended = prompt + [out1, out2, 10, 11, 12]
gen = model.generate(list(extended))
resumed_out = [next(gen) for _ in range(3)]
# compare against fresh generation (no cache) of the same prompt
model._cached_tokens = []
gen = model.generate(list(extended))
fresh_out = [next(gen) for _ in range(3)]
self.assertEqual(fresh_out, resumed_out)
def test_temperature_zero_is_greedy(self):
"""Temperature 0 (or near 0) should produce deterministic output."""
model = Transformer(TEST_CONFIG)
tokens = list(range(1, 6))
results = [list(zip(range(5), model.generate(list(tokens)))) for _ in range(3)]
# all runs should produce the same tokens
self.assertEqual(results[0], results[1])
self.assertEqual(results[1], results[2])
def test_temperature_high_produces_variety(self):
"""High temperature should produce different outputs across runs."""
model = Transformer(TEST_CONFIG)
tokens = list(range(1, 6))
runs = set()
for _ in range(5):
gen = model.generate(list(tokens), temperature=2.0)
out = tuple(next(gen) for _ in range(10))
runs.add(out)
# with temperature=2.0, we should see at least 2 distinct outputs across 5 runs
self.assertGreater(len(runs), 1, "high temperature should produce varied outputs")
def test_temperature_passed_to_forward(self):
"""Temperature from generate should be passed through to __call__."""
model = Transformer(TEST_CONFIG)
captured_temps = []
def mock_call(self, tokens, start_pos, temperature, **kwargs):
captured_temps.append(float(temperature.item()))
return Tensor([[42]])
with patch.object(Transformer, '__call__', mock_call):
gen = model.generate([1, 2, 3], temperature=0.6)
next(gen)
self.assertAlmostEqual(captured_temps[-1], 0.6, places=5)
if __name__ == '__main__':
unittest.main()

29
artifacts/package_sources/tinygrad/test/unit/test_masked_tensor.py Normal file
View File

@@ -0,0 +1,29 @@
import unittest
from tinygrad.tensor import Tensor
class TestMaskedTensor(unittest.TestCase):
def test_mul_masked(self):
a = Tensor([1,1,1,1,1])
b = Tensor([1,1]).pad(((0,3),))
c = a*b
assert c.shape == a.shape
ret = c.data()
assert ret.tolist() == [1.0, 1.0, 0.0, 0.0, 0.0]
def test_mul_both_masked(self):
a = Tensor([1,1]).pad(((0,3),))
b = Tensor([1,1]).pad(((0,3),))
c = a*b
assert c.shape == a.shape
ret = c.data()
assert ret.tolist() == [1.0, 1.0, 0.0, 0.0, 0.0]
def test_add_masked(self):
a = Tensor([1,1]).pad(((0,2),))
b = Tensor([1,1]).pad(((0,2),))
c = a+b
ret = c.data()
assert ret.tolist() == [2.0, 2.0, 0.0, 0.0]
if __name__ == '__main__':
unittest.main()

46
artifacts/package_sources/tinygrad/test/unit/test_metal_graph.py Normal file
View File

@@ -0,0 +1,46 @@
import unittest
from unittest.mock import MagicMock
from tinygrad import Device
from tinygrad.uop.ops import Ops, UOp
from tinygrad.dtype import dtypes
@unittest.skipUnless(Device.DEFAULT == "METAL", "Metal device required to run")
class TestMetalGraph(unittest.TestCase):
def setUp(self):
from tinygrad.runtime.graph.metal import MetalGraph
self.MetalGraph = MetalGraph
self.dev = Device[Device.DEFAULT]
def metal_buf(self, offset):
buf = MagicMock()
if offset > 0:
buf.op = Ops.SLICE
src = MagicMock()
src.dtype = dtypes.uint8
buf.src = (src, UOp.const(offset))
buf.dtype = dtypes.uint8
else:
buf.op = Ops.BUFFER
buf.device = Device.DEFAULT
return buf
def call(self, *bufs):
c = MagicMock()
c.src = (MagicMock(op=Ops.PROGRAM),) + tuple(bufs)
return c
def test_supports_uop_normal_offset(self):
assert self.MetalGraph.supports_uop([self.dev], self.call(self.metal_buf(0), self.metal_buf(100), self.metal_buf(0xFFFFFFFF))) is True
def test_supports_uop_overflow_offset(self):
assert self.MetalGraph.supports_uop([self.dev], self.call(self.metal_buf(0), self.metal_buf(0x100000000))) is False
def test_supports_uop_nonmetal_buf(self):
# non-SLICE ops should not be checked for offset
buf = MagicMock()
buf.op = Ops.BUFFER
buf.device = Device.DEFAULT
self.MetalGraph.supports_uop([self.dev], self.call(buf))
if __name__ == "__main__":
unittest.main()

994
artifacts/package_sources/tinygrad/test/unit/test_multitensor.py Normal file
View File

@@ -0,0 +1,994 @@
import unittest, numpy as np
from tinygrad import Tensor, Variable, Context, Device, TinyJit, GlobalCounters, dtypes, UOp, nn, getenv
from tinygrad.nn.state import get_parameters, get_state_dict
from tinygrad.uop.ops import Ops
from test.helpers import not_support_multi_device, needs_second_gpu, slow, assert_kernel_count, KernelCountException
from hypothesis import given, strategies as strat, settings
settings.register_profile("my_profile", max_examples=200, deadline=None, derandomize=getenv("DERANDOMIZE_CI", False))
settings.load_profile("my_profile")
d0 = f"{Device.DEFAULT}:0"
d1 = f"{Device.DEFAULT}:1"
d2 = f"{Device.DEFAULT}:2"
d3 = f"{Device.DEFAULT}:3"
d4 = f"{Device.DEFAULT}:4"
d5 = f"{Device.DEFAULT}:5"
devices_2 = (d1, d2)
devices_3 = (d1, d2, d3)
devices_4 = (d1, d2, d3, d4)
N = 128
@unittest.skipIf(not_support_multi_device(), "no multi")
class TestMultiTensor(unittest.TestCase):
@needs_second_gpu
def setUp(self): pass
def test_arange_shrink(self):
x = Tensor.arange(4)
self.assertEqual(x.shard(devices_2, 0).realize().shrink(((2, 4),)).tolist(), [2, 3])
self.assertEqual(x.shard(devices_2, 0).realize().shrink(((0, 2),)).tolist(), [0, 1])
def test_shard_like(self):
X = Tensor.ones(256).shard(devices_2, 0)
Y = Tensor.zeros(256).shard_like(X)
self.assertEqual(Y.device, X.device)
self.assertEqual(Y.uop.axis, 0)
# also test with axis=None
X2 = Tensor.ones(256).shard(devices_2, axis=None)
Y2 = Tensor.zeros(256).shard_like(X2)
self.assertEqual(Y2.device, X2.device)
self.assertEqual(Y2.uop.axis, None)
# test with single device
X3 = Tensor.ones(256)
Y3 = Tensor.zeros(256).shard_like(X3)
self.assertEqual(Y3.device, X3.device)
# cannot shard_like multi unless it's a no-op
X4 = Tensor.ones(256).shard(devices_2, 0)
Y4 = Tensor.ones(256).shard(devices_2, 0).shard_like(X4)
self.assertEqual(Y4.device, X4.device)
self.assertEqual(Y4.uop.axis, 0)
with self.assertRaises(RuntimeError):
Tensor.ones(256).shard(devices_2, None).shard_like(X4)
def _test_shard_op(self, op, out, n=4):
t = Tensor.ones(n).contiguous().realize().shard(devices_2, 0)
r = op(t).realize()
#assert t.uop.is_realized, "shard didn't realize"
self.assertEqual(r.tolist(), out)
def test_shard_reshape(self): self._test_shard_op(lambda t:t.reshape(2, 2), [[1.,1.],[1.,1.]])
def test_shard_elementwise(self): self._test_shard_op(lambda t:(t+t).reshape(2, 2), [[2.,2.],[2.,2.]])
def test_alu_deviceless_const(self):
s = Tensor([1.0, 2, 3, 4]).shard((f"{Device.DEFAULT}:0", f"{Device.DEFAULT}:1"), axis=0)
np.testing.assert_equal((s + Tensor(UOp.const(1.0).cast(dtypes.float))).numpy(), [2, 3, 4, 5])
np.testing.assert_equal((s + Tensor(UOp.const(1.0).cast(dtypes.float)).reshape((1,)).expand((4,))).numpy(), [2, 3, 4, 5])
def test_add_rank_expand_shard(self):
# a sharded src keeps its own rank under implicit broadcast, its shard axis right-aligns into the output
a = Tensor([1.,2.,3.,4.]).shard(devices_2, 0)
b = Tensor([[10.,20.,30.,40.]]).shard(devices_2, None)
self.assertEqual((a+b).uop.axis, 1)
np.testing.assert_equal((a+b).numpy(), [[11.,22.,33.,44.]])
def test_shard_reduce(self):
self._test_shard_op(lambda t:t.reshape(2, 3).sum(axis=1), [3.,3.], n=6)
self._test_shard_op(lambda t:t.reshape(2, 3).sum(axis=0), [2.,2.,2.], n=6)
def test_shard_not_multiple(self):
X = Tensor.ones(256).contiguous().realize()
with self.assertRaises(RuntimeError):
X.shard_(devices_3, 0)
def test_shard_reshape_cross_boundary(self):
X = Tensor.ones(5, 4).contiguous().realize().shard(devices_2, 1)
with self.assertRaises(RuntimeError): X.reshape(10, 2).uop.axis
def test_tensor_from_multi(self):
X = Tensor([1, 2], dtype=dtypes.int).shard_(devices_2, 0)
Y = Tensor(X.uop)
self.assertEqual(Y.device, devices_2)
np.testing.assert_equal(X.numpy(), Y.numpy())
Z = Tensor(X.uop, dtype=dtypes.float)
self.assertEqual(Z.dtype, dtypes.float)
np.testing.assert_equal(Z.numpy(), [1.0, 2.0])
def test_sharded_arange(self):
sharded_arange = Tensor.arange(1000).clone().shard(devices_2, 0)
sharded_arange.realize()
np.testing.assert_equal(sharded_arange.numpy(), np.arange(1000))
def test_shard_plus_one_sum(self):
X = Tensor.ones(256).contiguous().realize()
X.shard_((d1, d2), 0)
(X + 1).sum().realize()
def test_shard_plus_one_sum_d0(self):
X = Tensor.ones(256).contiguous().realize()
X.shard_((d0, d2), 0)
(X + 1).sum().realize()
def _test_simple_add_axis(self, shard_x, shard_w):
X = Tensor.ones(256).contiguous().realize()
W = Tensor.ones(256).contiguous().realize()
X.shard_((d1, d2), shard_x)
W.shard_((d1, d2), shard_w)
O = X + W
np.testing.assert_allclose(O.numpy(), 2)
def test_simple_add(self): return self._test_simple_add_axis(None, None)
def test_simple_add_X(self): return self._test_simple_add_axis(0, None)
def test_simple_add_W(self): return self._test_simple_add_axis(None, 0)
def test_simple_add_XW(self): return self._test_simple_add_axis(0, 0)
@given(strat.sampled_from((0, 1, None)), strat.sampled_from((0, 2)))
def test_allreduce_shard_ring_sum(self, axis, use_ring):
t = Tensor([1, 2, 3, 4]).reshape(2, 2)
with Context(RING=use_ring):
np.testing.assert_equal(t.shard(devices_2, axis=axis).sum().item(), 10)
def test_allreduce_cast_half(self, assign=False, kernel_count=8):
devices = tuple(f"{Device.DEFAULT}:{i}" for i in range(2))
a_src = Tensor.arange(2*3, dtype=dtypes.half).reshape(2, 3).clone().realize()
b_src = Tensor.arange(2*3, dtype=dtypes.half).reshape(2, 3).clone().realize()
a = a_src.shard(devices, axis=0).realize()
b = b_src.shard(devices, axis=0).realize()
# assigning creates a copy of the output before allreduce
if assign:
tst = Tensor.empty_like(b)
tst.assign(a + b)
else:
tst = a + b
tst = tst.float().sum(0)
GlobalCounters.reset()
with Context(ALLREDUCE_CAST=1, RING=0, ALL2ALL=0):
tst.realize()
assert_kernel_count(kernel_count)
np.testing.assert_allclose(tst.numpy(), (a_src.numpy()+b_src.numpy()).sum(0))
def test_allreduce_cast_half_assign(self): self.test_allreduce_cast_half(assign=True, kernel_count=10)
def test_multiple_to_single_device(self):
kernel_counts = {}
for ring in (0, 2):
GlobalCounters.reset()
with Context(RING=ring, SCACHE=0):
t = Tensor.arange(32).clone().shard(devices_4, 0).to(Device.DEFAULT)
t.realize()
kernel_counts[ring] = GlobalCounters.kernel_count
self.assertEqual(t.device, Device.DEFAULT)
np.testing.assert_equal(t.numpy(), np.arange(32))
self.assertEqual(kernel_counts[0], kernel_counts[2])
def test_to_single_device_gather_memory(self):
nrows, ncols = 64, 1024
nbytes = nrows*ncols*4
for devs in (devices_2, devices_4):
ndev = len(devs)
for axis in (0, 1):
sh = Tensor.arange(nrows*ncols).reshape(nrows, ncols).clone().shard(devs, axis).realize()
kernels, mem = {}, {}
for ring in (0, 2):
GlobalCounters.reset()
with Context(RING=ring, SCACHE=0):
t = sh.to(Device.DEFAULT)
t.realize()
kernels[ring], mem[ring] = GlobalCounters.kernel_count, GlobalCounters.global_mem
self.assertEqual(t.device, Device.DEFAULT)
np.testing.assert_equal(t.numpy(), np.arange(nrows*ncols).reshape(nrows, ncols))
self.assertEqual(kernels[0], kernels[2])
self.assertEqual(mem[0], mem[2])
self.assertLess(kernels[0], 2*ndev)
self.assertLessEqual(mem[0], 4*nbytes)
def test_multitensor_jit_input_reduce_shard_axis(self):
@TinyJit
def f(x): return x.sum(0).realize()
for _ in range(5):
tt = Tensor.ones(2, 64).contiguous().realize().shard((d1,d2), 0).realize()
out = f(tt)
np.testing.assert_allclose(out.numpy(), np.full(64, 2.0))
def _test_matmul_shard_axis(self, shard_x, shard_w, device):
X = Tensor.kaiming_uniform(N, N).realize()
W = Tensor.kaiming_uniform(N, N).realize()
Xs = X.shard(device, shard_x)
Ws = W.shard(device, shard_w)
O = (Xs@Ws)
with np.errstate(all='ignore'):
np.testing.assert_allclose(X.numpy() @ W.numpy(), O.to(Device.DEFAULT).numpy(), atol=1e-5)
def _test_double_matmul_shard_axis(self, shard_x, shard_w, device):
X = Tensor.kaiming_uniform(N, N).realize()
W1 = Tensor.kaiming_uniform(N, N).realize()
W2 = Tensor.kaiming_uniform(N, N).realize()
Xs = X.shard(device, shard_x)
W1s = W1.shard(device, shard_w)
W2s = W2.shard(device, shard_w)
O = (Xs@W1s)@W2s
with np.errstate(all='ignore'):
np.testing.assert_allclose((X.numpy() @ W1.numpy()) @ W2.numpy(), O.to(Device.DEFAULT).numpy(), atol=1e-5)
def test_matmul_shard_none(self): return self._test_matmul_shard_axis(None, None, devices_2)
def test_matmul_shard_X_0(self): return self._test_matmul_shard_axis(0, None, devices_2)
def test_matmul_shard_X_1(self): return self._test_matmul_shard_axis(1, None, devices_2)
def test_matmul_shard_W_0(self): return self._test_matmul_shard_axis(None, 0, devices_2)
def test_matmul_shard_W_1(self): return self._test_matmul_shard_axis(None, 1, devices_2)
def test_matmul_shard_0_0(self): return self._test_matmul_shard_axis(0, 0, devices_2)
def test_matmul_shard_0_1(self): return self._test_matmul_shard_axis(0, 1, devices_2)
def test_matmul_shard_1_0(self): return self._test_matmul_shard_axis(1, 0, devices_2)
def test_matmul_shard_1_1(self): return self._test_matmul_shard_axis(1, 1, devices_2)
def test_double_matmul_shard_X_0(self): return self._test_double_matmul_shard_axis(0, None, devices_2)
def test_double_matmul_shard_X_1(self): return self._test_double_matmul_shard_axis(1, None, devices_2)
def test_double_matmul_shard_W_0(self): return self._test_double_matmul_shard_axis(None, 0, devices_2)
def test_double_matmul_shard_W_1(self): return self._test_double_matmul_shard_axis(None, 1, devices_2)
def test_conv_data_shard(self):
conv = nn.Conv2d(3, 16, 3, bias=False)
for p in get_parameters(conv): p.shard_(devices_2)
fake_image = Tensor.rand((2, 3, 32, 32)).shard(devices_2, axis=0)
out = conv(fake_image)
out.numpy()
def test_conv_bias_data_shard(self):
conv = nn.Conv2d(3, 16, 3)
for p in get_parameters(conv): p.shard_(devices_2)
fake_image = Tensor.rand((2, 3, 32, 32)).shard(devices_2, axis=0)
out = conv(fake_image)
out.numpy()
def test_backprop_conv(self):
with Context(TRAINING=1):
conv = nn.Conv2d(3, 16, 3)
for p in get_parameters(conv): p.shard_(devices_2)
optim = nn.optim.Adam(get_parameters(conv))
fake_image = Tensor.rand((2, 3, 32, 32)).shard(devices_2, axis=0)
out = conv(fake_image)
optim.zero_grad()
out.mean().backward()
#for p in get_parameters(conv): p.grad.realize()
optim.step()
out.numpy()
def test_backprop_conv_wino(self):
with Context(WINO=1): self.test_backprop_conv()
def test_backward_sum(self):
x = Tensor([[1.,2,3,4], [5,6,7,8]]).shard(devices_2, axis=0)
w = Tensor([1.,2,3,4]).shard(devices_2)
out = x * w
out.mean().backward()
tst = w.grad.numpy()
np.testing.assert_allclose(tst, [0.75, 1., 1.25, 1.5])
def test_lr_scheduler_OneCycleLR(self):
from extra.lr_scheduler import OneCycleLR
conv = nn.Conv2d(3, 16, 3)
for p in get_parameters(conv): p.shard_(devices_2)
optim = nn.optim.SGD(get_parameters(conv))
lr_sched = OneCycleLR(optim, max_lr=0.1, pct_start=0.1, div_factor=100, final_div_factor=0.1, total_steps=10)
lr_sched.step()
def test_embedding(self):
B, T, embed_size, vocab_size = 4, 10, 20, 28
layer = nn.Embedding(vocab_size, embed_size)
x = Tensor(np.random.randint(0, vocab_size, (B, T), dtype=np.int32))
z = layer(x)
layer_sharded = nn.Embedding(vocab_size, embed_size)
layer_sharded.weight.replace(layer.weight.shard(devices_2, axis=1)).realize()
x_sharded = x.shard(devices_2, axis=None)
z_shard = layer_sharded(x_sharded)
np.testing.assert_allclose(z.numpy(), z_shard.numpy(), atol=1e-6, rtol=1e-6)
def test_embedding_backward(self, shard_weight_axis=None):
B, T, embed_size, vocab_size = 4, 10, 20, 28
layer = nn.Embedding(vocab_size, embed_size)
x = Tensor(np.random.randint(0, vocab_size, (B, T), dtype=np.int32))
z = layer(x)
z.sum().backward()
grad = layer.weight.grad.numpy()
layer_sharded = nn.Embedding(vocab_size, embed_size)
layer_sharded.weight.replace(layer.weight.shard(devices_2, axis=shard_weight_axis)).realize()
x_sharded = x.shard(devices_2, axis=None)
z_shard = layer_sharded(x_sharded)
z_shard.sum().backward()
grad_shard = layer_sharded.weight.grad.numpy()
np.testing.assert_allclose(grad, grad_shard, atol=1e-6, rtol=1e-6)
def test_embedding_backward_shard_weight(self): self.test_embedding_backward(shard_weight_axis=1)
def test_rmsnorm(self):
B, T, embed_size = 4, 10, 20
norm = nn.RMSNorm(embed_size)
x = Tensor.rand((B, T, embed_size)).contiguous().realize()
y = norm(x)
# for norm layers, the correct way to shard weights is duplication
norm_sharded = nn.RMSNorm(embed_size)
norm_sharded.weight.shard_(devices_2, axis=None).realize()
# if x is being sharded, then all-reduce is involved
x_sharded = x.shard(devices_2, axis=2).realize()
y_shard = norm_sharded(x_sharded).realize()
np.testing.assert_allclose(y.numpy(), y_shard.numpy(), atol=1e-6, rtol=1e-6)
# if x is being duplicated, then the operations remain inside each GPU
# which is the common case
x_sharded = x.shard(devices_2, axis=None).realize()
y_shard = norm_sharded(x_sharded).realize()
np.testing.assert_allclose(y.numpy(), y_shard.numpy(), atol=1e-6, rtol=1e-6)
def test_sdpa_causal_shard_batch(self):
B, H, T, D = 4, 2, 10, 16
q = Tensor.rand(B, H, T, D)
k = Tensor.rand(B, H, T, D)
v = Tensor.rand(B, H, T, D)
q_shard = q.shard(devices_2, axis=0)
k_shard = k.shard(devices_2, axis=0)
v_shard = v.shard(devices_2, axis=0)
Tensor.realize(q, k, v, q_shard, k_shard, v_shard)
y = Tensor.scaled_dot_product_attention(q, k, v, is_causal=True).realize()
y_shard = Tensor.scaled_dot_product_attention(q_shard, k_shard, v_shard, is_causal=True).realize()
np.testing.assert_allclose(y_shard.numpy(), y.numpy(), atol=1e-6, rtol=1e-6)
# NOTE: this is failing on LLVM CI, no idea why. Works locally.
@slow
def test_data_parallel_resnet(self):
from extra.models.resnet import ResNet18
fake_image = Tensor.rand((2, 3, 224//16, 224//16))
fake_image_sharded = fake_image.shard(devices_2, axis=0)
m = ResNet18()
m.load_from_pretrained()
real_output = m(fake_image).log_softmax().numpy()
for p in get_parameters(m): p.shard_(devices_2).realize()
GlobalCounters.reset()
shard_output = m(fake_image_sharded).log_softmax().realize()
shard_output_np = shard_output.numpy()
np.testing.assert_allclose(real_output, shard_output_np, atol=1e-6, rtol=1e-6)
def test_multi_tensor_jit_param(self):
@TinyJit
def jf(a, b) -> Tensor:
return (a + b).realize()
for _ in range(5):
a = Tensor.ones(256).contiguous().realize()
b = Tensor.ones(256).contiguous().realize()
a.shard_(devices_2)
b.shard_(devices_2)
c = jf(a, b)
np.testing.assert_allclose(c.numpy(), a.numpy()+b.numpy(), atol=1e-4, rtol=1e-5)
assert jf.captured is not None
def test_multi_tensor_jit_body(self):
@TinyJit
def jf() -> Tensor:
a = Tensor.ones(256).contiguous().realize()
b = Tensor.ones(256).contiguous().realize()
a.shard_(devices_2)
b.shard_(devices_2)
return (a + b).realize()
for _ in range(5):
r = jf()
np.testing.assert_allclose(r.numpy(), np.ones(256)+np.ones(256), atol=1e-4, rtol=1e-5)
assert jf.captured is not None
def test_multitensor_jit_in_list(self):
# test MULTI tensor inside a list container - exercises the container unpacking + MULTI unpacking
@TinyJit
def f(a, arr): return (a + arr[0]).realize()
for i in range(5):
a = Tensor.full((4,), i).contiguous().realize().shard(devices_2, 0).realize()
b = Tensor.ones(4).contiguous().realize().shard(devices_2, 0).realize()
out = f(a, [b])
np.testing.assert_allclose(out.numpy(), np.full(4, i) + np.ones(4), atol=1e-4, rtol=1e-5)
def test_multitensor_jit_multiple_inputs(self):
# test multiple MULTI tensors as inputs - each gets unpacked to component UOps
@TinyJit
def f(a, b, c): return (a + b + c).realize()
for i in range(5):
a = Tensor.full((4,), i).contiguous().realize().shard(devices_2, 0).realize()
b = Tensor.full((4,), i*2).contiguous().realize().shard(devices_2, 0).realize()
c = Tensor.ones(4).contiguous().realize().shard(devices_2, 0).realize()
out = f(a, b, c)
np.testing.assert_allclose(out.numpy(), np.full(4, i) + np.full(4, i*2) + np.ones(4), atol=1e-4, rtol=1e-5)
def test_multitensor_jit_different_sharding(self):
# test MULTI tensors with different sharding - one sharded on axis 0, one broadcast (axis=None)
@TinyJit
def f(a, b): return (a + b).realize()
for i in range(5):
a = Tensor.full((4, 4), i).contiguous().realize().shard(devices_2, 0).realize()
b = Tensor.full((4, 4), i*2).contiguous().realize().shard(devices_2, None).realize()
out = f(a, b)
np.testing.assert_allclose(out.numpy(), np.full((4, 4), i) + np.full((4, 4), i*2), atol=1e-4, rtol=1e-5)
def test_bn_ast_on_devices(self):
t = Tensor.empty((16, 64, 112, 112)).shard(devices_4, axis=0)
bn = nn.BatchNorm2d(64)
for p in get_parameters(bn): p.shard_(devices_4).realize()
out = bn(t)
scheds = [call for call in out.schedule_linear().src if call.src[0].op is not Ops.COPY and set(call.device) <= set(devices_4)]
self.assertEqual(set(scheds[0].device), set(devices_4), "should have ast on each shard device")
self.assertEqual(len(set(s.src[0] for s in scheds)), 1)
def test_flip(self):
rng = Tensor.rand((10, 10, 10))
t0 = rng.shard(devices_2, axis=1)
out = t0.flip(0) + 1
self.assertTrue((rng.flip(0)+1).allclose(out.to(rng.device)).item())
@unittest.skip("flaky")
def test_reshape_on_axis(self):
t0 = Tensor.rand((26, 15, 7)).shard(devices_3, axis=1)
# test split and rejoin to the right
t1 = t0.reshape((26, 3, 5, 7))
t2 = t0.reshape((26, 3, 35))
t3 = t1.reshape((26, 15, 7))
t4 = t2.reshape((26, 105,))
for t in [t0, t1, t2, t3, t4]:
assert t.uop.axis == 1
np.testing.assert_allclose(t.numpy().flatten(), t0.numpy().flatten())
# test shape-one axis
t5 = t4.reshape((26, 1, 105))
assert t5.uop.axis == 2
np.testing.assert_allclose(t.numpy().flatten(), t5.numpy().flatten())
# test split and rejoin to the right and reshape to the left
t5 = t0.reshape((2, 13, 3, 5, 7))
t6 = t0.reshape((13, 2, 3, 7, 5))
t7 = t0.reshape((1, 13, 2, 3, 1, 7, 5))
assert t5.uop.axis == 2
assert t6.uop.axis == 2
assert t7.uop.axis == 3
np.testing.assert_allclose(t5.numpy().flatten(), t0.numpy().flatten())
np.testing.assert_allclose(t6.numpy().flatten(), t0.numpy().flatten())
np.testing.assert_allclose(t7.numpy().flatten(), t0.numpy().flatten())
# test no left join
with self.assertRaises((AssertionError, ValueError)):
t0.reshape((26*15,7)).contiguous().schedule_linear()
# it doesn't work like this anymore
# NOTE: this never failed in assign_multi, it failed tensor spec because MULTI was never pushed in the graph
@unittest.skip("this test is broken")
def test_mlb_assign_change_axis(self):
t_none = Tensor.zeros((16, 16)).shard(devices_2).contiguous().realize()
t_zero = Tensor.ones((16, 16)).shard(devices_2, axis=0)
with self.assertRaises(RuntimeError):
# don't allow assigns that change axes
t_none.assign(t_zero)
t_none.schedule_linear()
def test_init_rand_with_multiple_devices_fail(self):
# init rand with multi device is not allowed
with self.assertRaises(ValueError):
Tensor.rand(256, device=devices_2)
def test_rand_like_on_shard(self, axis=None):
t = Tensor.empty((16, 16)).shard(devices_2, axis=axis)
t2 = Tensor.rand_like(t)
self.assertEqual(t.shape, t2.shape)
self.assertEqual(t.device, t2.device)
self.assertEqual(t.dtype, t2.dtype)
self.assertEqual(t.uop.axis, t2.uop.axis)
t2.realize()
def test_rand_like_on_shard_axis(self): self.test_rand_like_on_shard(0)
def test_rand_like_from_alu(self):
a = Tensor.ones(4, 4).shard(devices_4, axis=0)
aa = a + a
self.assertEqual(aa.device, devices_4)
self.assertEqual(aa.uop.axis, 0)
raa = aa.rand_like()
self.assertEqual(raa.device, devices_4)
self.assertEqual(raa.uop.axis, 0)
b = Tensor.empty(4, 4).shard(devices_4, axis=None)
ab = a + b
self.assertEqual(ab.device, devices_4)
self.assertEqual(ab.uop.axis, 0)
rab = ab.rand_like()
self.assertEqual(rab.device, devices_4)
self.assertEqual(rab.uop.axis, 0)
def test_rand_like_none_shard(self):
t = Tensor.empty((16, 16)).shard(devices_2)
t2 = Tensor.rand_like(t)
self.assertEqual(t.shape, t2.shape)
self.assertEqual(t.device, t2.device)
self.assertEqual(t.dtype, t2.dtype)
self.assertEqual(t.uop.axis, t2.uop.axis)
def test_rand_like_arg_dtype(self):
t = Tensor.empty((16, 16), dtype=dtypes.int32).shard(devices_2, axis=1)
t2 = Tensor.rand_like(t, dtype=dtypes.float32)
self.assertEqual(t.dtype, dtypes.int32)
self.assertEqual(t2.dtype, dtypes.float32)
def test_rand_like_arg_device(self):
# axis=None
t = Tensor.empty((16, 16)).shard((d1, d2), axis=None)
with self.assertRaises(RuntimeError):
Tensor.rand_like(t, device=(d3, d4))
# axis=1
t = Tensor.empty((16, 16)).shard((d1, d2), axis=1)
with self.assertRaises(RuntimeError):
Tensor.rand_like(t, device=(d3, d4))
def test_full_like_on_shard(self, axis=None):
t = Tensor.empty((16, 16)).shard(devices_2, axis=axis)
t2 = Tensor.full_like(t, 1.0)
self.assertEqual(t.shape, t2.shape)
self.assertEqual(t.device, t2.device)
self.assertEqual(t.dtype, t2.dtype)
self.assertEqual(t.uop.axis, t2.uop.axis)
t2.realize()
def test_full_like_on_shard_axis(self): self.test_full_like_on_shard(0)
def test_dropout_on_shard(self):
with Context(TRAINING=1):
X = Tensor.ones(256).to(devices_2)
output = X.dropout(0.5).numpy()
unique, counts = np.unique(output, return_counts=True)
assert set(unique) == {0, 2}, unique
assert 96 < counts[0] < 160, counts[0]
def test_dropout_on_shard_axis(self):
with Context(TRAINING=1):
X = Tensor.ones(512).shard(devices_2, axis=0)
output = X.dropout(0.5).numpy()
unique, counts = np.unique(output, return_counts=True)
assert set(unique) == {0, 2}, unique
assert 192 < counts[0] < 320, counts[0]
@unittest.skip("TODO: this requires forced_realize to be deleted.")
def test_shard_memory(self):
devices = (d0, d1, d2, d3)
t = Tensor.zeros(16, 16).contiguous()
t.shard_(devices, axis=0).realize()
assert all([lb is lb.base and lb.realized.base.size == 4 * 16 for lb in t.uop.src])
def test_clone(self):
for axis in (None, 0):
t = Tensor.arange(16).reshape(4, 4).clone().shard(devices_2, axis=axis).contiguous().realize()
t_clone = t.clone().realize()
self.assertEqual(t_clone.device, t.device)
self.assertEqual(t_clone.uop.axis, axis)
self.assertEqual(t_clone.tolist(), t.tolist())
t_clone += 1
self.assertNotEqual(t_clone.tolist(), t.tolist())
@unittest.skip("RANGEIFY doesn't support multi const folding")
def test_multi_const_folding(self):
with Context(TRACK_MATCH_STATS=0):
a = Tensor.arange(3).clone().realize()
zeros = Tensor.zeros(3).realize()
b = a.to(devices_2)*zeros.to(devices_2)
sched = b.schedule_linear().src
if len(sched) != 0: raise KernelCountException(0, len(sched))
self.assertListEqual(b.tolist(), [0, 0, 0])
@unittest.skipIf(not_support_multi_device(), "no multi")
class TestShrinkMultiTensorShardedAxis(unittest.TestCase):
@needs_second_gpu
def setUp(self): pass
# shrink a multitensor on sharded axis
def test_shrink_bad_args(self):
t = Tensor.arange(64).reshape(8, 8).clone().realize()
t.shard_([f"{Device.DEFAULT}:{i}" for i in range(4)], axis=0)
with self.assertRaises(RuntimeError):
# sharded axis shrink on non-device boundry is not allowed
a = t.shrink(((0, 3), (0, 8))).contiguous()
a.schedule_linear()
a = t.shrink(((0, 2), (2, 4)))
assert a.shape == (2, 2)
ref = Tensor.arange(64).reshape(8, 8).shrink(((0, 2), (2, 4)))
np.testing.assert_equal(a.numpy(), ref.numpy())
a = t.shrink(((0, 2), (0, 8))).contiguous()
a.schedule_linear()
assert a.shape == (2, 8)
p = a.pad(((0, 6), (0, 0))).contiguous()
p.schedule_linear()
assert p.shape == (8, 8)
@given(strat.sampled_from([dtypes.float, dtypes.int, dtypes.int64, dtypes.int16]))
def test_ops(self, dtype):
if dtype not in Device[Device.DEFAULT].renderer.supported_dtypes(): return
t = Tensor.arange(64).reshape(8, 8).clone().realize()
t.shard_([f"{Device.DEFAULT}:{i}" for i in range(4)], axis=0)
for i in range(2):
print(f"{i=}")
a = t.shrink(((0+2*i,2+2*i),None))
b = Tensor(t.numpy()[0+2*i:2+2*i])
assert a.shape == b.shape == (2, 8)
np.testing.assert_allclose(a.numpy(), b.numpy())
# cast
np.testing.assert_allclose(a.float().numpy(), b.float().numpy())
# elementwise
np.testing.assert_allclose(a.exp().numpy(), b.exp().numpy(), rtol=1e-7, atol=1e-3)
np.testing.assert_allclose(a.reciprocal().numpy(), b.reciprocal().numpy(), rtol=1e-7, atol=1e-3)
np.testing.assert_allclose(a.pow(-0.5).numpy(), b.pow(-0.5).numpy(), rtol=1e-7, atol=1e-3)
np.testing.assert_allclose((a+a).numpy(), (b+b).numpy(), rtol=1e-7, atol=1e-3)
np.testing.assert_equal((a+1).numpy(), (b+1).numpy())
np.testing.assert_equal((1+a).numpy(), (1+b).numpy())
np.testing.assert_allclose((a.bool().where(a+a, a)).numpy(), (b.bool().where(b+b, b)).numpy(), rtol=1e-7, atol=1e-3)
np.testing.assert_allclose((a.bool().where(1, 0)).numpy(), (b.bool().where(1, 0)).numpy(), rtol=1e-7, atol=1e-3)
# reduce
np.testing.assert_allclose(a.max().numpy(), b.max().numpy(), rtol=1e-7, atol=1e-3)
np.testing.assert_allclose(a.sum().numpy(), b.sum().numpy(), rtol=1e-7, atol=1e-3)
np.testing.assert_allclose(a.mean().numpy(), b.mean().numpy(), rtol=1e-7, atol=1e-3)
np.testing.assert_allclose(a.max(0).numpy(), b.max(0).numpy(), rtol=1e-7, atol=1e-3)
np.testing.assert_allclose(a.sum(0).numpy(), b.sum(0).numpy(), rtol=1e-7, atol=1e-3)
np.testing.assert_allclose(a.mean(0).numpy(), b.mean(0).numpy(), rtol=1e-7, atol=1e-3)
np.testing.assert_allclose(a.max(1).numpy(), b.max(1).numpy(), rtol=1e-7, atol=1e-3)
np.testing.assert_allclose(a.sum(1).numpy(), b.sum(1).numpy(), rtol=1e-7, atol=1e-3)
np.testing.assert_allclose(a.mean(1).numpy(), b.mean(1).numpy(), rtol=1e-7, atol=1e-3)
# pad it back
np.testing.assert_allclose(a.pad(((2*i, 2*(4-i-1)), None)).numpy(), b.pad(((2*i, 2*(4-i-1)), None)).numpy(), rtol=1e-7, atol=1e-3)
# other movement
np.testing.assert_allclose(a.pad((None, (1, 1))).numpy(), b.pad((None, (1, 1))).numpy(), rtol=1e-7, atol=1e-3)
np.testing.assert_allclose(a.shrink((None, (1, 3))).numpy(), b.shrink((None, (1, 3))).numpy(), rtol=1e-7, atol=1e-3)
np.testing.assert_allclose(a.permute((1, 0)).numpy(), b.permute((1, 0)).numpy(), rtol=1e-7, atol=1e-3)
np.testing.assert_allclose(a.reshape((2, 2, 4)).numpy(), b.reshape((2, 2, 4)).numpy(), rtol=1e-7, atol=1e-3)
np.testing.assert_allclose(a.reshape((2, 1, 8)).expand((2, 5, 8)).numpy(), b.reshape((2, 1, 8)).expand((2, 5, 8)).numpy(), rtol=1e-7, atol=1e-3)
np.testing.assert_allclose(a.flip(-1).numpy(), b.flip(-1).numpy(), rtol=1e-7, atol=1e-3)
def test_add_two_partitions(self):
t = Tensor.arange(64).reshape(8, 8).clone().realize()
t.shard_([f"{Device.DEFAULT}:{i}" for i in range(4)], axis=0)
a = t.shrink(((2, 4), None))
b = t.shrink(((6, 8), None))
na = t.numpy()[2:4]
nb = t.numpy()[6:8]
np.testing.assert_equal(a.numpy(), na)
np.testing.assert_equal(b.numpy(), nb)
np.testing.assert_equal((a+b).numpy(), na+nb)
c = a.pad(((2, 4), None)) + b.pad(((6, 0), None))
c.realize()
expected = np.concatenate([np.zeros_like(t.numpy()[0:2]), na, np.zeros_like(t.numpy()[4:6]), nb])
np.testing.assert_equal(c.numpy(), expected)
def test_add_different_tensors(self):
devices = [f"{Device.DEFAULT}:{i}" for i in range(4)]
x = Tensor.arange(64).reshape(8, 8).clone().realize().shard(devices, axis=0)
to_add = []
for i in range(len(devices)):
to_add.append((Tensor.ones(2, 8) * i).shard(devices))
added:list[Tensor] = []
for bound, a in zip(x.uop.bounds, to_add):
added.append(x[bound[0]:bound[1]] + a)
output = added[0].cat(*added[1:])
expected = np.arange(64).reshape((8,8)) + np.array([[0,0,1,1,2,2,3,3] for _ in range(8)]).T
np.testing.assert_allclose(output.numpy(), expected)
@unittest.skipIf(not_support_multi_device(), "no multi")
class TestBatchNorm(unittest.TestCase):
@needs_second_gpu
def setUp(self): pass
def test_unsynced_backprop_conv_bn(self):
with Context(TRAINING=1):
from extra.lr_scheduler import OneCycleLR
convs = [nn.Conv2d(3, 16, 3), nn.Conv2d(3, 16, 3)]
bns = [nn.BatchNorm2d(16), nn.BatchNorm2d(16)]
for p in get_parameters(convs + bns):
p.shard_((d1, d2))
optim = nn.optim.Adam(get_parameters(convs + bns))
lr_sched = OneCycleLR(optim, max_lr=0.1, pct_start=0.1, div_factor=100, final_div_factor=0.1, total_steps=10)
lr_sched.step()
fake_image = Tensor.rand((8, 3, 32, 32)).shard((d1, d2), axis=0)
f1 = fake_image.shrink(((0, 4), None, None, None))
f2 = fake_image.shrink(((4, 8), None, None, None))
out1 = bns[0](convs[0](f1))
out2 = bns[1](convs[1](f2))
out = out1.cat(out2)
optim.zero_grad()
out.mean().backward()
optim.step()
out.numpy()
def test_unsynced_backprop_standalone_bn(self):
from extra.lr_scheduler import OneCycleLR
GPUS = (d1, d2)
class BatchNorm:
def __init__(self, num_features):
self.bns:list[nn.BatchNorm2d] = []
for _ in GPUS:
bn = nn.BatchNorm2d(num_features, track_running_stats=False, eps=1e-12, momentum=0.85, affine=True)
self.bns.append(bn)
def __call__(self, x:Tensor):
bn_ts = []
each = x.shape[0]//len(self.bns)
for i, bn in enumerate(self.bns):
xi = x.shrink(((each*(i), each*(i+1)), None, None, None))
bni = bn(xi)
bn_ts.append(bni)
return bn_ts[0].cat(*bn_ts[1:])
with Context(TRAINING=1):
conv = nn.Conv2d(3, 16, 3)
bn = BatchNorm(16)
for p in get_parameters([conv, bn]):
p.shard_(GPUS)
optim = nn.optim.Adam(get_parameters([conv, bn]))
lr_sched = OneCycleLR(optim, max_lr=0.1, pct_start=0.1, div_factor=100, final_div_factor=0.1, total_steps=10)
lr_sched.step()
fake_image = Tensor.rand((8, 3, 32, 32)).shard(GPUS, axis=0)
out = bn(conv(fake_image))
optim.zero_grad()
out.mean().backward()
optim.step()
def test_unsynced_backprop_sync_weights(self):
from extra.lr_scheduler import OneCycleLR
from examples.hlb_cifar10 import UnsyncedBatchNorm
GPUS = (d1, d2)
with Context(TRAINING=1):
conv = nn.Conv2d(3, 16, 3)
bn = UnsyncedBatchNorm(16, num_devices=len(GPUS))
for k, p in get_state_dict([conv, bn]).items():
if 'running_mean' in k or 'running_var' in k:
p.shard_(GPUS, axis=0)
else:
p.to_(GPUS)
optim = nn.optim.Adam(get_parameters([conv, bn]))
lr_sched = OneCycleLR(optim, max_lr=0.1, pct_start=0.1, div_factor=100, final_div_factor=0.1, total_steps=10)
lr_sched.step()
fake_image = Tensor.rand((8, 3, 32, 32)).shard(GPUS, axis=0)
out = bn(conv(fake_image))
optim.zero_grad()
out.mean().backward()
optim.step()
@given(strat.sampled_from((False, True)))
def test_batchnorm(self, is_training):
devices = [f"{Device.DEFAULT}:{i}" for i in range(4)]
x = Tensor.arange(4096).reshape(8, 8, 8, 8).clone().realize().shard(devices, axis=0)
with Context(TRAINING=is_training):
bns = []
for _ in range(len(devices)):
bn = nn.BatchNorm2d(8)
for p in get_parameters(bn):
p.shard_(devices)
bns.append(bn)
bn_ts = []
for bound, bn in zip(x.uop.bounds, bns):
bni = bn(x[bound[0]:bound[1]])
bn_ts.append(bni)
bn_ts[0].cat(*bn_ts[1:]).numpy()
def test_synced_vs_unsynced_bn(self):
from examples.hlb_cifar10 import UnsyncedBatchNorm
from tinygrad.nn import BatchNorm2d
devices = [f"{Device.DEFAULT}:{i}" for i in range(4)]
x = Tensor.ones(8, 8, 8, 8).contiguous().realize().shard(devices, axis=0)
with Context(TRAINING=1):
synced_bn = BatchNorm2d(8)
unsynced_bn = UnsyncedBatchNorm(8, num_devices=len(devices))
for p in get_parameters(synced_bn):
p.shard_(devices)
for k, p in get_state_dict(unsynced_bn).items():
if 'running_mean' in k or 'running_var' in k:
p.shard_(devices, axis=0)
else:
p.to_(devices)
synced_out = synced_bn(x)
synced_si = list(synced_out.schedule_linear().src)
unsynced_out = unsynced_bn(x)
unsynced_si = list(unsynced_out.schedule_linear().src)
# TODO: test synced / unsynced batchnorm cross device kernel and copies
assert synced_si
assert unsynced_si
@unittest.skipIf(not_support_multi_device(), "need multi")
class TestMultiFromUnrenderable(unittest.TestCase):
@needs_second_gpu
def test_from_npy(self):
t = Tensor(np.arange(100, dtype=np.uint32))
ll = t.shard((d0, d1), axis=0) + 1
np.testing.assert_equal(ll.numpy(), np.arange(100)+1)
@unittest.skipIf(not_support_multi_device(), "need multi")
class TestMultiAssign(unittest.TestCase):
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(2))
@needs_second_gpu
def setUp(self): pass
def test_multi_assign_realized(self):
out = Tensor.zeros(4).shard(self.device, 0).contiguous().realize()
ones = Tensor.ones(4).shard(self.device, 0).contiguous().realize()
out.assign(ones).realize()
self.assertListEqual(out.tolist(), [1,1,1,1])
def test_multi_assign_unrealized(self):
out = Tensor.zeros(4).contiguous().realize().shard(self.device, 0)
ones = Tensor.ones(4).shard(self.device, 0).contiguous().realize()
out.assign(ones).realize()
self.assertListEqual(out.tolist(), [1,1,1,1])
def test_multi_assign_both_unrealized(self):
out = Tensor.zeros(4).contiguous().realize().shard(self.device, 0)
ones = Tensor.ones(4).contiguous().realize().shard(self.device, 0)
out.assign(ones).realize()
self.assertListEqual(out.tolist(), [1,1,1,1])
def test_multi_assign_scalar(self):
out = Tensor.ones(4).shard(self.device, 0).contiguous().realize()
out.assign(0).realize()
self.assertListEqual(out.tolist(), [0,0,0,0])
def test_multi_assign_const_like(self):
out = Tensor.ones(4).shard(self.device, 0).contiguous().realize()
out.assign(out.const_like(7)).realize()
self.assertListEqual(out.tolist(), [7,7,7,7])
def test_multi_assign_piece(self):
out = Tensor.zeros(4,4).shard(self.device, 0).contiguous().realize()
ones = Tensor.ones(4,1).shard(self.device, 0).contiguous().realize()
out[:, 2:3].assign(ones).realize()
self.assertListEqual(out.tolist(), [[0,0,1,0], [0,0,1,0], [0,0,1,0], [0,0,1,0]])
def test_multi_assign_piece_noncontig(self):
out = Tensor.zeros(4,4).contiguous().realize().shard(self.device, 0).realize()
ones = Tensor.ones(4,1).shard(self.device, 0).contiguous().realize()
out[:, 2:3].assign(ones).realize()
self.assertListEqual(out.tolist(), [[0,0,1,0], [0,0,1,0], [0,0,1,0], [0,0,1,0]])
@unittest.expectedFailure
def test_multi_assign_piece_unrealized(self):
out = Tensor.zeros(4,4).contiguous().realize().shard(self.device, 0)
ones = Tensor.ones(4,1).shard(self.device, 0).contiguous().realize()
out[:, 2:3].assign(ones).realize()
self.assertListEqual(out.tolist(), [[0,0,1,0], [0,0,1,0], [0,0,1,0], [0,0,1,0]])
def test_multi_assign_var_offset(self):
out = Tensor.zeros(4,4).contiguous().realize().shard(self.device, 0).realize()
ones = Tensor.ones(4,1).shard(self.device, 0).contiguous().realize()
vi = Variable("i", 0, 3).bind(2)
out[:, vi:vi+1].assign(ones).realize()
self.assertListEqual(out.tolist(), [[0,0,1,0], [0,0,1,0], [0,0,1,0], [0,0,1,0]])
def test_multi_assign_var_offset_jit_none(self): self.test_multi_assign_var_offset_jit(None)
def test_multi_assign_var_offset_jit(self, shard_axis=0):
out = Tensor.zeros(4,6).contiguous().realize().shard(self.device, shard_axis).realize()
ones = Tensor.ones(4,1).shard(self.device, shard_axis).contiguous().realize()
@TinyJit
def f(out:Tensor, vi):
out[:, vi:vi+1].assign(ones).realize()
ones.assign(ones+1).realize()
vi = Variable("i", 0, 5)
for i in range(1,5):
GlobalCounters.reset()
f(out, vi.bind(i))
self.assertListEqual(out.tolist(), [[0,1,2,3,4,0]]*4)
@unittest.skipIf(not_support_multi_device(), "need multi")
class TestMultiSetitem(unittest.TestCase):
device = tuple(f"{Device.DEFAULT}:{i}" for i in range(4))
@needs_second_gpu
def setUp(self): pass
def _t(self, axis): return Tensor.arange(16).clone().realize().shard(self.device, axis=axis)
def test_setitem_scalar_axis0(self):
t = self._t(0)
t[1] = 99
self.assertListEqual(t.tolist(), [0,99,2,3,4,5,6,7,8,9,10,11,12,13,14,15])
def test_setitem_scalar_axis_none(self):
t = self._t(None)
t[1] = 99
self.assertListEqual(t.tolist(), [0,99,2,3,4,5,6,7,8,9,10,11,12,13,14,15])
def test_setitem_slice_cross_shard(self):
t = self._t(0)
t[2:6] = 99
self.assertListEqual(t.tolist(), [0,1,99,99,99,99,6,7,8,9,10,11,12,13,14,15])
def test_setitem_full_slice(self):
t = self._t(0)
t[:] = 42
self.assertListEqual(t.tolist(), [42]*16)
def test_setitem_stride(self):
t = self._t(0)
t[::4] = 0
self.assertListEqual(t.tolist(), [0,1,2,3,0,5,6,7,0,9,10,11,0,13,14,15])
def test_setitem_single_shard(self):
t = self._t(0)
t[13] = 99
self.assertListEqual(t.tolist(), [0,1,2,3,4,5,6,7,8,9,10,11,12,99,14,15])
def test_setitem_tensor_value_replicated(self):
t = self._t(0)
t[2:6] = Tensor([90, 91, 92, 93]).shard(self.device)
self.assertListEqual(t.tolist(), [0,1,90,91,92,93,6,7,8,9,10,11,12,13,14,15])
def test_setitem_tensor_value_sharded_aligned(self):
t = self._t(0)
t[::4] = Tensor([90, 91, 92, 93]).shard(self.device, axis=0)
self.assertListEqual(t.tolist(), [90,1,2,3,91,5,6,7,92,9,10,11,93,13,14,15])
def helper_test_shard_op(shps, fxn, atol=1e-6, rtol=1e-3):
for shp in shps:
single_in = Tensor.randn(shp)
multi_in = single_in.shard(devices_2, axis=0)
single_out = fxn(single_in).numpy()
multi_out = fxn(multi_in).numpy()
try:
assert single_out.shape == multi_out.shape, f"shape mismatch: single={single_out.shape} | multi={multi_out.shape}"
assert single_out.dtype == multi_out.dtype, f"dtype mismatch: single={single_out.dtype} | multi={multi_out.dtype}"
np.testing.assert_allclose(single_out, multi_out, atol=atol, rtol=rtol)
except Exception as e:
raise Exception(f"Failed shape {single_out.shape}: {e}")
@unittest.skipIf(not_support_multi_device(), "no multi")
class TestTensorOps(unittest.TestCase):
@needs_second_gpu
def test_interpolate(self):
helper_test_shard_op([(4,16,16),(4,24,24)], lambda x: Tensor.interpolate(x, (19,19)))
@needs_second_gpu
def test_bitcast(self):
helper_test_shard_op([(256,), (256,)], lambda x: x.bitcast(dtypes.int))
if __name__ == '__main__':
unittest.main()

17
artifacts/package_sources/tinygrad/test/unit/test_objc.py Normal file
View File

@@ -0,0 +1,17 @@
import sys, unittest
class TestObjCMetaSpec(unittest.TestCase):
@unittest.skipUnless(sys.platform == "darwin", "objc runtime only on macOS")
def test_classmethods_are_classmethods(self):
from tinygrad.runtime.support.objc import Spec, id_
#_classmethods_ must include classmethod descriptors
class ObjCTest(Spec):
_methods_ = [("foo", id_, [])]
_classmethods_ = [("bar", id_, [])]
self.assertNotIsInstance(ObjCTest.__dict__["foo"], classmethod)
self.assertIsInstance(ObjCTest.__dict__["bar"], classmethod)
if __name__ == "__main__":
unittest.main()

200
artifacts/package_sources/tinygrad/test/unit/test_randomness.py Normal file
View File

@@ -0,0 +1,200 @@
import unittest, math, torch
import numpy as np
from functools import partial
from tinygrad import nn, dtypes, Tensor, Device, TinyJit, Variable
from tinygrad.helpers import OSX
# https://gist.github.com/devries/11405101
def ksprob(a):
fac, total, termbf = 2.0, 0.0, 0.0
a2 = -2.0 * a * a
for j in range(1, 101):
term = fac * math.exp(a2 * j * j)
total += term
if math.fabs(term) <= 0.001 * termbf or math.fabs(term) <= 1e-8 * total:
return total
fac = -fac
termbf = math.fabs(term)
return 1.0
def kstest(l1, l2):
n1, n2 = len(l1), len(l2)
l1.sort()
l2.sort()
j1, j2, d, fn1, fn2 = 0, 0, 0.0, 0.0, 0.0
while j1 < n1 and j2 < n2:
d1, d2 = l1[j1], l2[j2]
if d1 <= d2:
fn1 = (float(j1) + 1.0) / float(n1)
j1 += 1
if d2 <= d1:
fn2 = (float(j2) + 1.0) / float(n2)
j2 += 1
dtemp = math.fabs(fn2 - fn1)
if dtemp > d:
d = dtemp
ne = float(n1 * n2) / float(n1 + n2)
nesq = math.sqrt(ne)
prob = ksprob((nesq + 0.12 + 0.11 / nesq) * d)
return prob
def equal_distribution(tiny_func, torch_func=None, numpy_func=None, shape=(40, 43), alpha=0.04):
Tensor.manual_seed(1337)
torch.manual_seed(1337)
np.random.seed(1337)
assert not (torch_func is None and numpy_func is None), "no function to compare with"
x1 = tiny_func(*shape).numpy().flatten()
x2 = tiny_func(shape).numpy().flatten()
if numpy_func is not None: y = numpy_func(shape).flatten()
if torch_func is not None: z = torch_func(shape).numpy().flatten()
return (numpy_func is None or (kstest(x1, y) >= alpha and kstest(x2, y) >= alpha)) and \
(torch_func is None or (kstest(x1, z) >= alpha and kstest(x2, z) >= alpha))
def normal_test(func, shape=(20, 45), alpha=0.05): return equal_distribution(func, numpy_func=lambda x: np.random.randn(*x), shape=shape, alpha=alpha)
class TestRandomness(unittest.TestCase):
def test_three_lazy_rands_realized_one_at_a_time_are_distinct(self):
Tensor.manual_seed(123)
r1, r2, r3 = [Tensor.rand(4) for _ in range(3)]
self.assertNotEqual(r1.tolist(), r2.tolist())
self.assertNotEqual(r2.tolist(), r3.tolist())
def test_randn(self):
self.assertEqual(Tensor.randn(3,3,dtype=dtypes.half).dtype, dtypes.half)
self.assertTrue(normal_test(Tensor.randn))
self.assertTrue(equal_distribution(Tensor.randn, torch.randn, lambda x: np.random.randn(*x)))
def test_randint(self):
self.assertFalse(normal_test(Tensor.randint))
self.assertTrue(equal_distribution(partial(Tensor.randint, low=-2, high=5),
numpy_func=lambda x: np.random.randint(low=-2, high=5, size=x)))
self.assertTrue(equal_distribution(partial(Tensor.randint, low=-2, high=5, dtype="int32"),
numpy_func=lambda x: np.random.randint(low=-2, high=5, size=x)))
self.assertTrue(Tensor.randint(1, device="CPU").device=="CPU")
# check types of args
with self.assertRaises(TypeError): Tensor.randint((3, 4), low=0.1, high=3)
with self.assertRaises(TypeError): Tensor.randint((3, 4), low=0, high=3.5)
with self.assertRaises(TypeError): Tensor.randint((3, 4), low=1, high=3, dtype="float")
with self.assertRaises(TypeError): Tensor.randint((3, 4), low=0, high=3, dtype=dtypes.float32)
# check low < high
with self.assertRaises(ValueError): Tensor.randint((3, 4), low=10, high=5)
with self.assertRaises(ValueError): Tensor.randint((3, 4), low=10, high=10)
np.testing.assert_array_equal(Tensor.randint(16, low=5, high=6).numpy(), 5)
def test_normal(self):
self.assertTrue(normal_test(Tensor.normal))
self.assertTrue(equal_distribution(Tensor.normal, lambda x: torch.nn.init.normal_(torch.empty(x), mean=0, std=1),
lambda x: np.random.normal(loc=0, scale=1, size=x)))
# check std >= 0
with self.assertRaises(ValueError): Tensor.normal((3, 4), mean=0, std=-1)
def test_uniform(self):
self.assertFalse(normal_test(Tensor.uniform))
self.assertTrue(equal_distribution(Tensor.uniform, lambda x: torch.nn.init.uniform_(torch.empty(x)), lambda x: np.random.uniform(size=x)))
self.assertTrue(equal_distribution(partial(Tensor.uniform, low=-100, high=100, dtype=dtypes.int32),
numpy_func=lambda x: np.random.randint(low=-100, high=100, size=x)))
# check low < high
with self.assertRaises(ValueError): Tensor.uniform((3, 4), low=5.0, high=3.0)
with self.assertRaises(ValueError): Tensor.uniform((3, 4), low=1.0, high=1.0)
def test_scaled_uniform(self):
self.assertFalse(normal_test(Tensor.scaled_uniform))
self.assertTrue(equal_distribution(Tensor.scaled_uniform, lambda x: torch.nn.init.uniform_(torch.empty(x), a=-1, b=1) / math.sqrt(math.prod(x)),
lambda x: np.random.uniform(-1, 1, size=x) / math.sqrt(math.prod(x))))
def test_glorot_uniform(self):
self.assertFalse(normal_test(Tensor.glorot_uniform))
self.assertTrue(equal_distribution(Tensor.glorot_uniform, lambda x: torch.nn.init.xavier_uniform_(torch.empty(x)),
lambda x: np.random.uniform(-1, 1, size=x) * math.sqrt(6 / (x[0] + math.prod(x[1:])))))
def test_kaiming_uniform(self):
for shape in [(32, 16, 3, 3), (20, 44), (5, 15, 35)]:
self.assertTrue(equal_distribution(Tensor.kaiming_uniform, lambda x: torch.nn.init.kaiming_uniform_(torch.empty(x)), shape=shape))
def test_kaiming_normal(self):
for shape in [(32, 16, 3, 3), (20, 44), (3, 15, 35)]:
self.assertTrue(equal_distribution(Tensor.kaiming_normal, lambda x: torch.nn.init.kaiming_normal_(torch.empty(x)), shape=shape))
def test_multinomial(self):
self.assertRaises(AssertionError, lambda: Tensor(2).multinomial(1, replacement=False))
self.assertRaises(AssertionError, lambda: Tensor([1, 9]).multinomial(0, replacement=False))
def _check_with_torch(w, num_samples, replacement):
tiny_res = Tensor(w).multinomial(num_samples, replacement=replacement).realize()
torch_res = torch.tensor(w).multinomial(num_samples, replacement=replacement)
self.assertEqual(tiny_res.shape, torch_res.shape)
if torch_res.ndim == 1:
tiny_res = tiny_res.unsqueeze(0)
torch_res = torch_res.unsqueeze(0)
for i in range(torch_res.shape[0]):
self.assertTrue(equal_distribution(lambda *_: tiny_res[i], lambda _: torch_res[i]))
_check_with_torch(w=[0.231, 0., 1., 0.5], num_samples=300, replacement=True)
_check_with_torch(w=[[0.2, 0.8]], num_samples=300, replacement=True) # 2D but only 1 row
_check_with_torch(w=[[0.453, 0., 1., 0.81], [0.1, 0.8, 0., 0.1]], num_samples=300, replacement=True)
# no-replacement
w = [0.1, 0.9]
self.assertRaises(AssertionError, lambda: Tensor(w).multinomial(100, replacement=False))
@TinyJit
def sample_one(): return Tensor(w).multinomial(1, replacement=False).realize()
tiny_samples = [sample_one().item() for _ in range(200)]
torch_samples = [torch.tensor(w).multinomial(1, replacement=False).item() for _ in range(200)]
self.assertTrue(equal_distribution(lambda *_: Tensor(tiny_samples), lambda _: torch.tensor(torch_samples)))
w = list(range(32))
s1 = Tensor(w).multinomial(5, replacement=False).numpy()
self.assertEqual(len(set(s1.tolist())), 5)
s2 = Tensor(w).multinomial(5, replacement=False).numpy()
self.assertFalse(np.array_equal(s1, s2))
full = Tensor(w).multinomial(len(w), replacement=False).numpy()
self.assertEqual(sorted(full.tolist()), w)
w = [0.1, 0.2, 0.3, 0.4]
@TinyJit
def sample_three(): return Tensor(w).multinomial(3, replacement=False).realize()
tiny_draws = np.array([sample_three().numpy() for _ in range(200)])
torch_draws = np.array([torch.tensor(w).multinomial(3, replacement=False).numpy() for _ in range(200)])
for pos in range(3):
self.assertTrue(equal_distribution(lambda *_: Tensor(tiny_draws[:, pos]), lambda _: torch.tensor(torch_draws[:, pos])))
@unittest.skip("this test is flaky")
def test_multinomial_counterexample(self):
tiny_res = Tensor([0.3, 0.6, 0.1]).multinomial(4000, replacement=True)
torch_res = torch.tensor([0.3, 0.6, 0.1]).multinomial(4000, replacement=True)
self.assertTrue(equal_distribution(lambda *_: tiny_res, lambda _: torch_res))
torch_res = torch.tensor([0.2, 0.7, 0.1]).multinomial(4000, replacement=True)
self.assertFalse(equal_distribution(lambda *_: tiny_res, lambda _: torch_res))
def test_conv2d_init(self):
params = (32, 64, (3,3))
assert equal_distribution(lambda *_: nn.Conv2d(*params).weight, lambda _: torch.nn.Conv2d(*params).weight.detach())
assert equal_distribution(lambda *_: nn.Conv2d(*params).bias, lambda _: torch.nn.Conv2d(*params).bias.detach())
def test_linear_init(self):
params = (64, 256)
assert equal_distribution(lambda *_: nn.Linear(*params).weight, lambda _: torch.nn.Linear(*params).weight.detach())
assert equal_distribution(lambda *_: nn.Linear(*params).bias, lambda _: torch.nn.Linear(*params).bias.detach())
def test_bn_init(self):
params = (64,)
assert equal_distribution(lambda *_: nn.BatchNorm2d(*params).weight, lambda _: torch.nn.BatchNorm2d(*params).weight.detach())
assert equal_distribution(lambda *_: nn.BatchNorm2d(*params).bias, lambda _: torch.nn.BatchNorm2d(*params).bias.detach())
# TODO: still fails with MAX_KERNEL_BUFFERS
@unittest.skipIf(Device.DEFAULT == "WEBGPU" and not OSX, "WEBGPU Vulkan can only run kernels with up to 10 buffers")
class TestSample(unittest.TestCase):
def test_sample(self):
X = Tensor.rand(1000, 50).realize()
BS = 16
idxs = np.random.randint(0, X.shape[0], size=(BS))
# this uncovered a bug with arg sort order
batch = [Variable(f'idx{i}', 0, X.shape[0]-1).bind(s) for i,s in enumerate(idxs.tolist())]
x = Tensor.cat(*[X.shrink(((batch[i], batch[i]+1), None)) for i in range(BS)])
print(idxs)
ret = x.numpy()
base = X.numpy()[idxs]
np.testing.assert_equal(ret, base)
if __name__ == "__main__":
unittest.main()

67
artifacts/package_sources/tinygrad/test/unit/test_realize_is_realize.py Normal file
View File

@@ -0,0 +1,67 @@
import tempfile, unittest
import numpy as np
from tinygrad import Tensor, Device, dtypes, Variable
class TestRealizeIsRealized(unittest.TestCase):
def test_list(self):
t = Tensor([1, 2, 3]).realize()
assert t.uop.is_realized
def test_rand(self):
t = Tensor.rand(4, 4).realize()
assert t.uop.is_realized
def test_contiguous(self):
t = Tensor.zeros(10).contiguous().realize()
assert t.uop.is_realized
def test_ones(self):
t = Tensor.ones(4, 4).realize()
assert t.uop.is_realized
def test_bytes(self):
t = Tensor(b'\x01\x02\x03').realize()
assert t.uop.is_realized
def test_numpy(self):
t = Tensor(np.array([1, 2, 3])).realize()
assert t.uop.is_realized
def test_multi(self):
d = Device.DEFAULT
t = Tensor.ones(8).contiguous().shard((d, d), axis=0).realize()
assert t.uop.src[0].is_realized
def test_empty(self):
t = Tensor.empty(4, 4).realize()
assert not t.uop.is_realized
def test_disk(self):
with tempfile.NamedTemporaryFile() as f:
f.write(b'\x00' * 16)
f.flush()
t = Tensor.empty(4, dtype=dtypes.float32, device=f"disk:{f.name}").realize()
assert not t.uop.is_realized
def test_assign(self):
t = Tensor([1, 2, 3])
t += 1
t.realize()
assert t.uop.is_realized
# TODO: these are not realized after .realize()
def test_const_not_realized(self):
t = Tensor(3.14).realize()
assert not t.uop.is_realized
def test_none_not_realized(self):
t = Tensor(None).realize()
assert not t.uop.is_realized
def test_variable_not_realized(self):
t = Tensor(Variable("v", 1, 10).bind(3)).realize()
assert not t.uop.is_realized
if __name__ == "__main__":
unittest.main()

166
artifacts/package_sources/tinygrad/test/unit/test_rearrange_einops.py Normal file
View File

@@ -0,0 +1,166 @@
# modified from
# https://github.com/arogozhnikov/einops/blob/master/tests/test_examples.py
# https://github.com/arogozhnikov/einops/blob/master/tests/test_ops.py
# https://github.com/arogozhnikov/einops/blob/master/tests/test_parsing.py
import numpy as np
import unittest
from tinygrad import Tensor
class test_rearrange_examples(unittest.TestCase):
def test_tensor_train_example_numpy(self):
# kept here just for a collection, only tested for numpy
# https://arxiv.org/pdf/1509.06569.pdf, (5)
x = Tensor.ones([3, 4, 5, 6])
rank = 4
# creating appropriate Gs
Gs = [Tensor.ones([d, d, rank, rank]) for d in x.shape]
Gs[0] = Gs[0][:, :, :1, :]
Gs[-1] = Gs[-1][:, :, :, :1]
# einsum way
y = x.reshape((1,) + x.shape)
for G in Gs:
# taking partial results left-to-right
# y = numpy.einsum('i j alpha beta, alpha i ... -> beta ... j', G, y)
y = Tensor(np.einsum("i j a b, a i ... -> b ... j", G.numpy(), y.numpy()))
y1 = y.reshape(-1)
# alternative way
y = x.reshape(-1)
for G in Gs:
i, j, alpha, beta = G.shape
y = y.rearrange("(i rest alpha) -> rest (alpha i)", alpha=alpha, i=i)
y = y @ G.rearrange("i j alpha beta -> (alpha i) (j beta)")
y = y.rearrange("rest (beta j) -> (beta rest j)", beta=beta, j=j)
y2 = y
assert np.allclose(y1.numpy(), y2.numpy())
# yet another way
y = x
for G in Gs:
i, j, alpha, beta = G.shape
y = y.rearrange("i ... (j alpha) -> ... j (alpha i)", alpha=alpha, i=i)
y = y @ G.rearrange("i j alpha beta -> (alpha i) (j beta)")
y3 = y.reshape(-1)
assert np.allclose(y1.numpy(), y3.numpy())
class test_rearrange_ops(unittest.TestCase):
def test_rearrange_ellipsis_ops(self):
identity_patterns = [
"...->...",
"a b c d e-> a b c d e",
"a b c d e ...-> ... a b c d e",
"a b c d e ...-> a ... b c d e",
"... a b c d e -> ... a b c d e",
"a ... e-> a ... e",
"a ... -> a ... ",
"a ... c d e -> a (...) c d e",
]
equivalent_rearrange_patterns = [
("a b c d e -> (a b) c d e", "a b ... -> (a b) ... "),
("a b c d e -> a b (c d) e", "... c d e -> ... (c d) e"),
("a b c d e -> a b c d e", "... -> ... "),
("a b c d e -> (a b c d e)", "... -> (...)"),
("a b c d e -> b (c d e) a", "a b ... -> b (...) a"),
("a b c d e -> b (a c d) e", "a b ... e -> b (a ...) e"),
]
xnp = np.arange(2 * 3 * 4 * 5 * 6, dtype=np.int32).reshape([2, 3, 4, 5, 6])
x = Tensor(xnp)
for pattern in identity_patterns:
assert np.array_equal(xnp, x.rearrange(pattern).numpy()), pattern
for pattern1, pattern2 in equivalent_rearrange_patterns:
assert np.array_equal(x.rearrange(pattern1).numpy(), x.rearrange(pattern2).numpy())
def test_rearrange_consistency(self):
shape = [1, 2, 3, 5, 7, 11]
xnp = np.arange(np.prod(shape), dtype=np.int32).reshape(shape)
x = Tensor(xnp)
for pattern in [
"a b c d e f -> a b c d e f",
"b a c d e f -> a b d e f c",
"a b c d e f -> f e d c b a",
"a b c d e f -> (f e) d (c b a)",
"a b c d e f -> (f e d c b a)",
]:
result = x.rearrange(pattern).numpy()
assert len(np.setdiff1d(xnp, result)) == 0
assert result.dtype == xnp.dtype
result = x.rearrange("a b c d e f -> a (b) (c d e) f").numpy()
assert np.array_equal(xnp.flatten(), result.flatten())
result = x.rearrange("a aa aa1 a1a1 aaaa a11 -> a aa aa1 a1a1 aaaa a11").numpy()
assert np.array_equal(xnp, result)
result1 = x.rearrange("a b c d e f -> f e d c b a").numpy()
result2 = x.rearrange("f e d c b a -> a b c d e f").numpy()
assert np.array_equal(result1, result2)
result = x.rearrange("a b c d e f -> (f d) c (e b) a").rearrange("(f d) c (e b) a -> a b c d e f", b=2, d=5).numpy()
assert np.array_equal(xnp, result)
sizes = dict(zip("abcdef", shape))
temp = x.rearrange("a b c d e f -> (f d) c (e b) a", **sizes)
result = temp.rearrange("(f d) c (e b) a -> a b c d e f", **sizes).numpy()
assert np.array_equal(xnp, result)
x2 = np.arange(2 * 3 * 4, dtype=np.int32).reshape([2, 3, 4])
result = Tensor(x2).rearrange("a b c -> b c a").numpy()
assert x2[1, 2, 3] == result[2, 3, 1]
assert x2[0, 1, 2] == result[1, 2, 0]
def test_rearrange_permutations(self):
# tests random permutation of axes against two independent numpy ways
for n_axes in range(1, 10):
x = np.arange(2**n_axes, dtype=np.int32).reshape([2] * n_axes)
permutation = np.random.permutation(n_axes)
left_expression = " ".join("i" + str(axis) for axis in range(n_axes))
right_expression = " ".join("i" + str(axis) for axis in permutation)
expression = left_expression + " -> " + right_expression
result = Tensor(x).rearrange(expression).numpy()
for pick in np.random.randint(0, 2, [10, n_axes]):
assert x[tuple(pick)] == result[tuple(pick[permutation])]
for n_axes in range(1, 10):
x = np.arange(2**n_axes, dtype=np.int32).reshape([2] * n_axes)
permutation = np.random.permutation(n_axes)
left_expression = " ".join("i" + str(axis) for axis in range(n_axes)[::-1])
right_expression = " ".join("i" + str(axis) for axis in permutation[::-1])
expression = left_expression + " -> " + right_expression
result = Tensor(x).rearrange(expression).numpy()
assert result.shape == x.shape
expected_result = np.zeros_like(x)
for original_axis, result_axis in enumerate(permutation):
expected_result |= ((x >> original_axis) & 1) << result_axis
assert np.array_equal(result, expected_result)
class test_rearrange_parsing(unittest.TestCase):
def test_unicode_ellipsis(self):
equivalent_rearrange_patterns = [
("a b … -> (a b) … ", "a b ... -> (a b) ... "),
("… c d e -> … (c d) e", "... c d e -> ... (c d) e"),
("… -> … ", "... -> ... "),
("… -> (…)", "... -> (...)"),
("a b … -> b (…) a", "a b ... -> b (...) a"),
("a b … e -> b (a …) e", "a b ... e -> b (a ...) e"),
]
xnp = np.arange(2 * 3 * 4 * 5 * 6, dtype=np.int32).reshape([2, 3, 4, 5, 6])
x = Tensor(xnp)
for pattern1, pattern2 in equivalent_rearrange_patterns:
assert np.array_equal(x.rearrange(pattern1).numpy(), x.rearrange(pattern2).numpy())
if __name__ == "__main__":
unittest.main()

69
artifacts/package_sources/tinygrad/test/unit/test_schedule_cache.py Normal file
View File

@@ -0,0 +1,69 @@
import unittest
import functools
from tinygrad import Tensor, Variable, UOp
from tinygrad.uop.ops import KernelInfo
from tinygrad.schedule import schedule_cache
def custom_set0_kernel(A:UOp, num:int) -> UOp:
return A[0].set(num).sink(arg=KernelInfo(f"custom_set0_{num}"))
class TestScheduleCache(unittest.TestCase):
def test_bound_variable_reuses_cache(self):
schedule_cache.clear()
v = Variable('v', 1, 100)
x = Tensor.ones(10).contiguous().realize()
# first run with v=5
t1 = (x + Tensor(v.bind(5))).sum()
self.assertEqual(t1.item(), 60.0)
cache_size_after_first = len(schedule_cache)
# second run with v=10 should reuse cache
t2 = (x + Tensor(v.bind(10))).sum()
self.assertEqual(t2.item(), 110.0)
self.assertEqual(len(schedule_cache), cache_size_after_first)
def test_custom_kernel(self):
for i in range(4):
a = Tensor.empty(1)
a = Tensor.custom_kernel(a, fxn=functools.partial(custom_set0_kernel, num=i))[0]
a.realize()
self.assertEqual(a.item(), i)
def test_same_custom_function_reuses_cache(self):
schedule_cache.clear()
fxn = functools.partial(custom_set0_kernel, num=10)
# first run
a = Tensor.empty(1)
a = Tensor.custom_kernel(a, fxn=fxn)[0]
a.realize()
self.assertEqual(a.item(), 10)
cache_size_after_first = len(schedule_cache)
# second run with same function should reuse cache
b = Tensor.empty(1)
b = Tensor.custom_kernel(b, fxn=fxn)[0]
b.realize()
self.assertEqual(b.item(), 10)
self.assertEqual(len(schedule_cache), cache_size_after_first)
def test_simple(self):
a = Tensor.ones(10).contiguous()
b = Tensor.ones(10).contiguous()
Tensor.realize(a, b)
# warm up
for _ in range(2):
num = (a.sum().contiguous()+b.sum().contiguous()).item()
print(num)
# confirm schedule cache doesn't grow
start_len_schedule_cache = len(schedule_cache)
for _ in range(3):
num = (a.sum().contiguous()+b.sum().contiguous()).item()
print(num)
self.assertEqual(len(schedule_cache), start_len_schedule_cache)
if __name__ == "__main__":
unittest.main()

143
artifacts/package_sources/tinygrad/test/unit/test_setitem_schedule.py Normal file
View File

@@ -0,0 +1,143 @@
import unittest
from tinygrad import Tensor, dtypes, GlobalCounters
from test.helpers import assert_kernel_count
class TestSetitemInto(unittest.TestCase):
def test_setitem_into_unrealized(self):
GlobalCounters.reset()
t = Tensor.arange(4, dtype=dtypes.int32).reshape(2, 2)
assert_kernel_count(0)
t[1] = 5
assert_kernel_count(0)
t.realize()
assert_kernel_count(0)
self.assertEqual(GlobalCounters.global_mem, 0)
self.assertListEqual(t.tolist(), [[0, 1], [5, 5]])
def test_setitem_into_unrealized_sliced_compute(self):
# base computation contains SHRINK from prior slicing (like QR decomposition pattern)
GlobalCounters.reset()
a = Tensor.arange(8, dtype=dtypes.int32).reshape(2, 4)
w = a[0] + a[1] # unrealized ADD with SHRINK in graph: [4, 6, 8, 10]
assert_kernel_count(0)
w[1] = 99
assert_kernel_count(0)
w.realize()
assert_kernel_count(0)
self.assertEqual(GlobalCounters.global_mem, 0)
self.assertListEqual(w.tolist(), [4, 99, 8, 10])
def test_setitem_into_empty(self):
GlobalCounters.reset()
t = Tensor.empty(4, dtype=dtypes.int32)
t[1] = 5
assert_kernel_count(0)
t.realize()
assert_kernel_count(1)
self.assertEqual(GlobalCounters.global_mem, 4)
t[1].realize()
t.realize()
assert_kernel_count(1)
self.assertEqual(t[1].item(), 5)
def test_setitem_into_empty_alu(self):
GlobalCounters.reset()
t = Tensor.empty(4, dtype=dtypes.int32) + 1
assert_kernel_count(0)
t[1] = 5
assert_kernel_count(0)
t.realize()
assert_kernel_count(1)
self.assertLessEqual(GlobalCounters.global_mem, 32)
t[1].realize()
t.realize()
assert_kernel_count(1)
self.assertEqual(t[1].item(), 5)
def test_setitem_into_tensor(self):
t = Tensor([1, 2, 3, 4], dtype=dtypes.int32).realize()
GlobalCounters.reset()
t[1] = 5
assert_kernel_count(0)
t[1].realize()
assert_kernel_count(1)
self.assertEqual(GlobalCounters.global_mem, 4)
t.realize()
assert_kernel_count(1)
self.assertListEqual(t.tolist(), [1, 5, 3, 4])
def test_setitem_into_tensor_alu(self):
t = Tensor([1, 2, 3, 4], dtype=dtypes.int32).realize() + 1
GlobalCounters.reset()
t[1] = 5
assert_kernel_count(0)
t[1].realize()
assert_kernel_count(1)
self.assertLessEqual(GlobalCounters.global_mem, 32)
t[1].realize()
t.realize()
assert_kernel_count(1)
self.assertListEqual(t.tolist(), [2, 5, 4, 5])
def test_setitem_into_const(self):
GlobalCounters.reset()
t = Tensor.ones(4, dtype=dtypes.int32, buffer=False)
t[1] = 5
assert_kernel_count(0)
t.realize()
assert_kernel_count(0)
self.assertEqual(GlobalCounters.global_mem, 0)
self.assertListEqual(t.tolist(), [1, 5, 1, 1])
def test_setitem_into_const_alu(self):
GlobalCounters.reset()
t = Tensor.ones(4, dtype=dtypes.int32, buffer=False) + 1
t[1] = 5
assert_kernel_count(0)
t.realize()
assert_kernel_count(0)
self.assertEqual(GlobalCounters.global_mem, 0)
self.assertListEqual(t.tolist(), [2, 5, 2, 2])
def test_setitem_into_arange(self):
# NOTE: arange has no real buffer, but assigning to it is fine
GlobalCounters.reset()
other = Tensor.arange(4, dtype=dtypes.int32)
t = Tensor.arange(4, dtype=dtypes.int32)
self.assertIs(other.uop, t.uop)
t[1] = 5
assert_kernel_count(0)
t.realize()
assert_kernel_count(0)
self.assertListEqual(t.tolist(), [0, 5, 2, 3])
def test_setitem_slice_const(self):
t = Tensor.zeros(100, dtype=dtypes.int32).contiguous().realize()
GlobalCounters.reset()
t[20:50] = 3
assert_kernel_count(0)
t.realize()
assert_kernel_count(1)
self.assertEqual(GlobalCounters.global_mem, 30*4) # 30 elements written
def test_setitem_slice_tensor(self):
t = Tensor.zeros(100, dtype=dtypes.int32).contiguous().realize()
v = Tensor.zeros(30, dtype=dtypes.int32).contiguous().realize()
GlobalCounters.reset()
t[20:50] = v
assert_kernel_count(0)
t.realize()
assert_kernel_count(1)
self.assertEqual(GlobalCounters.global_mem, 30*4*2) # 30 read + 30 written
def test_setitem_full(self):
t = Tensor.zeros(100, dtype=dtypes.int32).contiguous().realize()
GlobalCounters.reset()
t[:] = 3
assert_kernel_count(0)
t.realize()
assert_kernel_count(1)
self.assertEqual(GlobalCounters.global_mem, 100*4) # full buffer written
if __name__ == '__main__':
unittest.main()

40
artifacts/package_sources/tinygrad/test/unit/test_shm_tensor.py Normal file
View File

@@ -0,0 +1,40 @@
import unittest
import multiprocessing.shared_memory as shared_memory
from tinygrad.helpers import WIN
from tinygrad import Tensor, Device
import numpy as np
class TestRawShmBuffer(unittest.TestCase):
@unittest.skipIf(WIN, "only fails on CI windows instance")
def test_e2e(self):
t = Tensor.randn(2, 2, 2).realize()
# copy to shm
shm_name = (s := shared_memory.SharedMemory(create=True, size=t.nbytes())).name
s.close()
t_shm = t.to(f"disk:shm:{shm_name}").realize()
# copy from shm
t2 = t_shm.to(Device.DEFAULT).realize()
assert np.allclose(t.numpy(), t2.numpy())
s.unlink()
@unittest.skip("big shared memory")
def test_e2e_big(self):
# bigger than this doesn't work on Linux, maybe this is a limit somewhere?
t = Tensor.randn(2048, 128, 8).realize()
# copy to shm
shm_name = (s := shared_memory.SharedMemory(create=True, size=t.nbytes())).name
s.close()
t_shm = t.to(f"disk:shm:{shm_name}").realize()
# copy from shm
t2 = t_shm.to(Device.DEFAULT).realize()
assert np.allclose(t.numpy(), t2.numpy())
s.unlink()
if __name__ == "__main__":
unittest.main()

13
artifacts/package_sources/tinygrad/test/unit/test_symbolic_tensor.py Normal file
View File

@@ -0,0 +1,13 @@
import unittest
from tinygrad import Variable
from tinygrad.tensor import Tensor
class TestSymbolicPad(unittest.TestCase):
def test_pad(self):
v = Variable("v", 1, 100).bind(5)
t = Tensor.ones(100)[:v].pad(((4, 0),))
t = t[:9]
assert t.tolist() == [0,0,0,0,1,1,1,1,1]
if __name__ == '__main__':
unittest.main()

28
artifacts/package_sources/tinygrad/test/unit/test_system_pci_scan_bus.py Normal file
View File

@@ -0,0 +1,28 @@
import sys
import pytest
@pytest.mark.skipif(sys.platform != "linux", reason="uses linux sysfs layout")
def test_pci_scan_bus_filters_vendor(monkeypatch):
import tinygrad.runtime.support.system as system
fake = {
"/sys/bus/pci/devices/0000:00:01.0/vendor": "0x1234",
"/sys/bus/pci/devices/0000:00:01.0/device": "0x1111",
"/sys/bus/pci/devices/0000:00:02.0/vendor": "0xabcd",
"/sys/bus/pci/devices/0000:00:02.0/device": "0x1111",
}
class FakeFileIOInterface:
def __init__(self, path, *args, **kwargs):
self.path = path
def listdir(self):
assert self.path == "/sys/bus/pci/devices"
return ["0000:00:01.0", "0000:00:02.0"]
def read(self, *args, **kwargs):
return fake[self.path]
monkeypatch.setattr(system, "FileIOInterface", FakeFileIOInterface)
assert system.System.pci_scan_bus(0x1234, devices=[(0xffff, [0x1111])]) == ["0000:00:01.0"]

141
artifacts/package_sources/tinygrad/test/unit/test_tar.py Normal file
View File

@@ -0,0 +1,141 @@
import unittest, tarfile, io, os, pathlib, tempfile
import numpy as np
from tinygrad import Tensor
from tinygrad.nn.state import tar_extract
class TestTarExtractFile(unittest.TestCase):
def setUp(self):
self.test_dir = tempfile.mkdtemp()
self.test_files = {
'file1.txt': b'Hello, World!',
'file2.bin': b'\x00\x01\x02\x03\x04',
'empty_file.txt': b''
}
self.tar_path = os.path.join(self.test_dir, 'test.tar')
with tarfile.open(self.tar_path, 'w') as tar:
for filename, content in self.test_files.items():
file_path = os.path.join(self.test_dir, filename)
with open(file_path, 'wb') as f:
f.write(content)
tar.add(file_path, arcname=filename)
# Create invalid tar file
self.invalid_tar_path = os.path.join(self.test_dir, 'invalid.tar')
with open(self.invalid_tar_path, 'wb') as f:
f.write(b'This is not a valid tar file')
def tearDown(self):
for filename in self.test_files:
os.remove(os.path.join(self.test_dir, filename))
os.remove(self.tar_path)
os.remove(self.invalid_tar_path)
os.rmdir(self.test_dir)
def test_tar_extract_returns_dict(self):
result = tar_extract(self.tar_path)
self.assertIsInstance(result, dict)
def test_tar_extract_correct_keys(self):
result = tar_extract(self.tar_path)
self.assertEqual(set(result.keys()), set(self.test_files.keys()))
def test_tar_extract_content_size(self):
result = tar_extract(self.tar_path)
for filename, content in self.test_files.items():
self.assertEqual(len(result[filename]), len(content))
def test_tar_extract_content_values(self):
result = tar_extract(self.tar_path)
for filename, content in self.test_files.items():
np.testing.assert_array_equal(result[filename].numpy(), np.frombuffer(content, dtype=np.uint8))
def test_tar_extract_empty_file(self):
result = tar_extract(self.tar_path)
self.assertEqual(len(result['empty_file.txt']), 0)
def test_tar_extract_non_existent_file(self):
with self.assertRaises(FileNotFoundError):
tar_extract('non_existent_file.tar')
def test_tar_extract_invalid_file(self):
with self.assertRaises(tarfile.ReadError):
tar_extract(self.invalid_tar_path)
class TestTarExtractPAX(unittest.TestCase):
tar_format = tarfile.PAX_FORMAT
max_link_len = 1000_000
test_files = {
'a/file1.txt': b'Hello, World!',
'a/b/file2.bin': b'\x00\x01\x02\x03\x04',
'empty_file.txt': b'',
'512file': b'a' * 512,
'long_file': b'some data' * 100,
'very' * 15 + '/' + 'very' * 15 + '_long_filename.txt': b'Hello, World!!',
'very' * 200 + '_long_filename.txt': b'Hello, World!!!',
}
def create_tar_tensor(self):
fobj = io.BytesIO()
test_dirs = set(os.path.dirname(k) for k in self.test_files.keys()).difference({ '' })
with tarfile.open(fileobj=fobj, mode='w', format=self.tar_format) as tar:
for dirname in test_dirs:
dir_info = tarfile.TarInfo(name=dirname)
dir_info.type = tarfile.DIRTYPE
tar.addfile(dir_info)
for filename, content in self.test_files.items():
file_info = tarfile.TarInfo(name=filename)
file_info.size = len(content)
tar.addfile(file_info, io.BytesIO(content))
if len(filename) < self.max_link_len:
link_info = tarfile.TarInfo(name=filename + '.lnk')
link_info.type = tarfile.SYMTYPE
link_info.linkname = filename
tar.addfile(link_info)
return Tensor(fobj.getvalue())
def test_tar_extract_returns_dict(self):
result = tar_extract(self.create_tar_tensor())
self.assertIsInstance(result, dict)
def test_tar_extract_correct_keys(self):
result = tar_extract(self.create_tar_tensor())
self.assertEqual(set(result.keys()), set(self.test_files.keys()))
def test_tar_extract_content_size(self):
result = tar_extract(self.create_tar_tensor())
for filename, content in self.test_files.items():
self.assertEqual(len(result[filename]), len(content))
def test_tar_extract_content_values(self):
result = tar_extract(self.create_tar_tensor())
for filename, content in self.test_files.items():
np.testing.assert_array_equal(result[filename].numpy(), np.frombuffer(content, dtype=np.uint8))
def test_tar_extract_empty_file(self):
result = tar_extract(self.create_tar_tensor())
self.assertEqual(len(result['empty_file.txt']), 0)
def test_tar_extract_non_existent_file(self):
with self.assertRaises(FileNotFoundError):
tar_extract(Tensor(pathlib.Path('non_existent_file.tar')))
def test_tar_extract_invalid_file(self):
with self.assertRaises(tarfile.ReadError):
tar_extract(Tensor(b'This is not a valid tar file'))
def test_tar_extract_invalid_file_long(self):
with self.assertRaises(tarfile.ReadError):
tar_extract(Tensor(b'This is not a valid tar file'*100))
class TestTarExtractUSTAR(TestTarExtractPAX):
tar_format = tarfile.USTAR_FORMAT
max_link_len = 100
test_files = {k: v for k, v in TestTarExtractPAX.test_files.items() if len(k) < 256}
class TestTarExtractGNU(TestTarExtractPAX):
tar_format = tarfile.GNU_FORMAT
if __name__ == '__main__':
unittest.main()

92
artifacts/package_sources/tinygrad/test/unit/test_tensor_data.py Normal file
View File

@@ -0,0 +1,92 @@
import unittest, struct
from tinygrad import Tensor, dtypes
from tinygrad.uop.ops import UOp
# format types: https://docs.python.org/3/library/struct.html
class TestTensorBytes(unittest.TestCase):
def test_bytes(self):
lst = Tensor(bytes(b"\xaa\xbb\xcc\xdd"))
assert lst.tolist() == [170, 187, 204, 221]
def test_float_bytes(self):
lst = Tensor(bytes(struct.pack("ff", 0.234, 0.8585)), dtype=dtypes.float32)
assert lst.shape == (2,)
assert abs(lst.tolist()[0] - 0.234) < 1e-6
assert abs(lst.tolist()[1] - 0.8585) < 1e-6
class TestTensorData(unittest.TestCase):
def test_data(self):
a = Tensor([1,2,3,4], dtype=dtypes.int32)
dat = a.data()
assert dat.itemsize == 4
assert list(dat) == [1,2,3,4]
assert dat.shape == (4,)
assert dat[0] == 1
assert dat[1] == 2
def test_data_empty(self):
a = Tensor([], dtype=dtypes.int32)
dat = a.data()
assert dat.itemsize == 4
assert list(dat) == []
assert dat.shape == (0,)
def test_data_empty_multi_dim(self):
a = Tensor([], dtype=dtypes.int32).reshape(0, 2)
dat = a.data()
assert dat.itemsize == 4
assert list(dat) == []
assert dat.shape == (0,)
def test_data_uint8(self):
a = Tensor([1,2,3,4], dtype=dtypes.uint8)
dat = a.data()
assert dat.format == "B"
assert dat.itemsize == 1
assert dat[0] == 1
assert dat[1] == 2
def test_data_nested(self):
a = Tensor([[1,2],[3,4]], dtype=dtypes.int32)
dat = a.data()
assert dat.format == "i"
assert dat.itemsize == 4
assert dat.tolist() == [[1, 2], [3, 4]]
assert dat.shape == (2,2)
assert dat[0, 0] == 1
assert dat[1, 1] == 4
def test_data_const(self):
a = Tensor(3, dtype=dtypes.int32)
dat = a.data()
assert dat.format == "i"
assert dat.itemsize == 4
assert dat.tolist() == 3
assert dat.shape == ()
def test_const_dtype_for_uop(self):
self.assertEqual(Tensor.const(UOp.const(1.0).cast(dtypes.float32), dtypes.int8).dtype, dtypes.int8)
self.assertEqual(Tensor.const(UOp.variable("x", 1, 10).bind(5), dtypes.int32).item(), 5)
def test_data_float32(self):
a = Tensor([[1,2.5],[3,4]], dtype=dtypes.float32)
dat = a.data()
assert dat.format == "f"
assert dat[0, 1] == 2.5
@unittest.skip("requires python 3.12")
def test_data_float16(self):
a = Tensor([[1,2.5],[3,4]], dtype=dtypes.float16)
dat = a.data()
assert dat.format == "e"
assert dat.shape == (2,2)
# NOTE: python can't deref float16
def test_tolist_empty_shapes(self):
for shape, expected in (((0,), []), ((2, 0), [[], []]), ((0, 2), []),
((2, 0, 3), [[], []]), ((2, 3, 0), [[[], [], []], [[], [], []]])):
self.assertEqual(Tensor.ones(*shape).tolist(), expected)
if __name__ == '__main__':
unittest.main()

21
artifacts/package_sources/tinygrad/test/unit/test_tensor_io.py Normal file
View File

@@ -0,0 +1,21 @@
import unittest
from tinygrad import Tensor
from tinygrad.nn.state import TensorIO
class TestTensorIO(unittest.TestCase):
def test_read(self):
data = b"Hello World!"
fobj = TensorIO(Tensor(data))
self.assertEqual(fobj.read(1), data[:1])
self.assertEqual(fobj.read(5), data[1:6])
self.assertEqual(fobj.read(100), data[6:])
self.assertEqual(fobj.read(100), b"")
def test_read_nolen(self):
data = b"Hello World!"
fobj = TensorIO(Tensor(data))
fobj.seek(2)
self.assertEqual(fobj.read(), data[2:])
if __name__ == '__main__':
unittest.main()

83
artifacts/package_sources/tinygrad/test/unit/test_tinyfs.py Normal file
View File

@@ -0,0 +1,83 @@
import json, math, os, socketserver, threading, unittest
import numpy as np
from tinygrad import Tensor, dtypes
from tinygrad.helpers import CHUNK_SIZE
from tinygrad.nn.state import fs_store, fs_load
from extra.tinyfs.fetch_file import hash_file, _python_hash_1mb
_chunks: dict[bytes, bytes] = {}
class _Handler(socketserver.StreamRequestHandler):
def handle(self):
while line := self.rfile.readline():
cmd = line.decode().strip()
if cmd == "INFO":
self.wfile.write(json.dumps({"node0": ["node0", f"127.0.0.1:{self.server.server_address[1]}"]}).encode() + b"\r\n")
elif cmd.startswith("STORE_IN"):
data = self.rfile.read(int(cmd.split()[1]))
hashes = bytearray()
for i in range(math.ceil(len(data) / CHUNK_SIZE)):
chunk = data[i*CHUNK_SIZE:(i+1)*CHUNK_SIZE].ljust(CHUNK_SIZE, b'\0')
h = _python_hash_1mb(chunk)
_chunks[h] = chunk
hashes.extend(h)
self.wfile.write(hashes)
elif cmd.startswith("LOAD_IN"):
hashes = self.rfile.read(int(cmd.split()[1]))
self.wfile.write(json.dumps(["node0"] * (len(hashes) // 16)).encode() + b"\r\n")
elif cmd.startswith("CHUNK_OUT"):
size = int(cmd.split()[1])
self.wfile.write(_chunks.get(self.rfile.read(16), bytes(size))[:size])
self.wfile.flush()
# regressed in 55d3a5def "preallocate all realized buffers"
class TestTinyFS(unittest.TestCase):
@classmethod
def setUpClass(cls):
_chunks.clear()
cls._server = socketserver.ThreadingTCPServer(('127.0.0.1', 0), _Handler)
cls._server.daemon_threads = True
threading.Thread(target=cls._server.serve_forever, daemon=True).start()
os.environ["TINYFS_ENDPOINT"] = f"127.0.0.1:{cls._server.server_address[1]}"
@classmethod
def tearDownClass(cls):
_chunks.clear()
os.environ.pop("TINYFS_ENDPOINT", None)
cls._server.shutdown()
cls._server.server_close()
def test_store(self):
h = fs_store(Tensor([1.0, 2.0, 3.0, 4.0])).realize()
self.assertEqual(h.shape, (16,))
self.assertEqual(h.dtype, dtypes.uint8)
def test_store_deterministic(self):
a = fs_store(Tensor([1.0, 2.0, 3.0, 4.0])).realize()
b = fs_store(Tensor([1.0, 2.0, 3.0, 4.0])).realize()
np.testing.assert_array_equal(a.numpy(), b.numpy())
def test_store_different_data(self):
a = fs_store(Tensor([1.0, 2.0, 3.0, 4.0])).realize()
b = fs_store(Tensor([5.0, 6.0, 7.0, 8.0])).realize()
self.assertNotEqual(a.tolist(), b.tolist())
def test_roundtrip_uint8(self):
arr = np.arange(256, dtype=np.uint8)
loaded = fs_load(fs_store(Tensor(arr)).realize(), len(arr)).to("CPU")
np.testing.assert_array_equal(loaded.numpy(), arr)
def test_roundtrip_multichunk_uint8(self):
arr = np.random.default_rng(42).integers(0, 256, size=CHUNK_SIZE + 1024, dtype=np.uint8)
loaded = fs_load(fs_store(Tensor(arr)).realize(), len(arr)).to("CPU")
np.testing.assert_array_equal(loaded.numpy(), arr)
def test_hash_matches_python_impl(self):
arr = np.arange(256, dtype=np.uint8)
h = fs_store(Tensor(arr)).realize()
# the hash from fs_store should match the pure-Python hash_file reference
padded = arr.tobytes().ljust(CHUNK_SIZE, b'\0')
self.assertEqual(h.data().tobytes(), hash_file(padded))
if __name__ == "__main__":
unittest.main()

38
artifacts/package_sources/tinygrad/test/unit/test_winograd.py Normal file
View File

@@ -0,0 +1,38 @@
import unittest, sys
import numpy as np
from tinygrad import Tensor, GlobalCounters, Context, nn
from tinygrad.helpers import WINO
@unittest.skipIf(sys.platform.startswith("win"), "flaky on Windows")
class TestWinogradClose(unittest.TestCase):
def test_close(self):
inp = Tensor.rand(1, 16, 16, 16)
conv = nn.Conv2d(16, 16, 3)
conv(inp).realize() # warmup
GlobalCounters.reset()
print("non winograd")
with Context(WINO=0):
cmp = conv(inp).realize() # warmup
GlobalCounters.reset()
print("winograd")
with Context(WINO=1):
test = conv(inp).realize()
np.testing.assert_allclose(cmp.numpy(), test.numpy(), atol=1e-5)
@unittest.skipIf(sys.platform.startswith("win"), "flaky on Windows")
class TestWinograd(unittest.TestCase):
def setUp(self):
self.old = WINO.value
WINO.value = 1
def tearDown(self):
WINO.value = self.old
def test_padded_conv2d(self):
# tests padding order in winograd
x,w = Tensor.rand(1,3,11,28).realize(), Tensor.rand(4,3,3,3).realize()
with Context(WINO=0): expected = Tensor.conv2d(x,w,padding=1).realize()
with Context(WINO=1): result = Tensor.conv2d(x,w,padding=1).realize()
np.testing.assert_allclose(result.numpy(), expected.numpy(), atol=1e-4)
if __name__ == '__main__':
unittest.main(verbosity=2)