forked from IQ.Lvbs/IQ.Pilot
IQ.Pilot Release Commit @ 0798119
This commit is contained in:
IQ.Lvbs history cleanup
0
tinygrad_repo/test/unit/__init__.py
Normal file
0
tinygrad_repo/test/unit/__init__.py
Normal file
52
tinygrad_repo/test/unit/test_allreduce.py
Normal file
52
tinygrad_repo/test/unit/test_allreduce.py
Normal file
@@ -0,0 +1,52 @@
|
||||
import unittest
|
||||
from tinygrad import Tensor, 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_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()
|
||||
1002
tinygrad_repo/test/unit/test_assign.py
Normal file
1002
tinygrad_repo/test/unit/test_assign.py
Normal file
File diff suppressed because it is too large
Load Diff
203
tinygrad_repo/test/unit/test_attention.py
Normal file
203
tinygrad_repo/test/unit/test_attention.py
Normal file
@@ -0,0 +1,203 @@
|
||||
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, "full_attention_interval":2,
|
||||
"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}")
|
||||
|
||||
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()
|
||||
351
tinygrad_repo/test/unit/test_call.py
Normal file
351
tinygrad_repo/test/unit/test_call.py
Normal file
@@ -0,0 +1,351 @@
|
||||
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_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 — unique consts must not leak through
|
||||
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
tinygrad_repo/test/unit/test_callify.py
Normal file
111
tinygrad_repo/test/unit/test_callify.py
Normal 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()
|
||||
133
tinygrad_repo/test/unit/test_conv.py
Normal file
133
tinygrad_repo/test/unit/test_conv.py
Normal file
@@ -0,0 +1,133 @@
|
||||
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_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-1)
|
||||
np.testing.assert_allclose(r1.numpy(), np.maximum(out.numpy(), 0), atol=1e-5)
|
||||
np.testing.assert_allclose(r2.numpy(), out.numpy() - 1, atol=1e-5)
|
||||
|
||||
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
tinygrad_repo/test/unit/test_cpu.py
Normal file
26
tinygrad_repo/test/unit/test_cpu.py
Normal 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[3].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()
|
||||
568
tinygrad_repo/test/unit/test_disk_tensor.py
Normal file
568
tinygrad_repo/test/unit/test_disk_tensor.py
Normal file
@@ -0,0 +1,568 @@
|
||||
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
|
||||
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)
|
||||
|
||||
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.bfloat16]: continue # not supported in numpy
|
||||
path = self.tmp(f"ones.{dtype}.safetensors")
|
||||
ones = Tensor(np.random.rand(10,10), dtype=dtype)
|
||||
safe_save(get_state_dict(ones), path)
|
||||
np.testing.assert_equal(ones.numpy(), list(safe_load(path).values())[0].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.safetensors"))
|
||||
|
||||
loaded = safe_load(self.tmp("dtypes.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.safetensors"))
|
||||
|
||||
loaded = safe_load(self.tmp("dtypes.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())
|
||||
|
||||
@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()
|
||||
242
tinygrad_repo/test/unit/test_dtype_spec.py
Normal file
242
tinygrad_repo/test/unit/test_dtype_spec.py
Normal file
@@ -0,0 +1,242 @@
|
||||
import unittest, math, subprocess
|
||||
from tinygrad.tensor import Tensor
|
||||
from tinygrad.dtype import dtypes, DType, DTYPES_DICT
|
||||
from tinygrad.device import Device
|
||||
from tinygrad.helpers import getenv, DEBUG, EMULATED_DTYPES
|
||||
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
|
||||
# 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 setUp(self):
|
||||
self.old_default_int, self.old_default_float = dtypes.default_int, dtypes.default_float
|
||||
def tearDown(self):
|
||||
dtypes.default_int, dtypes.default_float = self.old_default_int, self.old_default_float
|
||||
|
||||
@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):
|
||||
dtypes.default_int, dtypes.default_float = default_int, default_float
|
||||
_assert_eq(Tensor(True), dtypes.bool, True)
|
||||
_assert_eq(Tensor(None), dtypes.default_float, [])
|
||||
_assert_eq(Tensor(2), dtypes.default_int, 2)
|
||||
_assert_eq(Tensor(2.34), dtypes.default_float, 2.34)
|
||||
_assert_eq(Tensor([]), dtypes.default_float, [])
|
||||
_assert_eq(Tensor([1]), dtypes.default_int, [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):
|
||||
dtypes.default_int, dtypes.default_float = default_int, 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):
|
||||
dtypes.default_int, dtypes.default_float = default_int, 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):
|
||||
dtypes.default_int, dtypes.default_float = default_int, 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 setUp(self):
|
||||
self.old_default_int, self.old_default_float = dtypes.default_int, dtypes.default_float
|
||||
def tearDown(self):
|
||||
dtypes.default_int, dtypes.default_float = self.old_default_int, self.old_default_float
|
||||
|
||||
@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):
|
||||
old_default_float = dtypes.default_float
|
||||
|
||||
for default_dtype in dtypes.floats:
|
||||
if default_dtype not in supported_dtypes: continue
|
||||
dtypes.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])
|
||||
|
||||
dtypes.default_float = old_default_float
|
||||
|
||||
@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)
|
||||
558
tinygrad_repo/test/unit/test_function.py
Normal file
558
tinygrad_repo/test/unit/test_function.py
Normal file
@@ -0,0 +1,558 @@
|
||||
import numpy as np
|
||||
import unittest
|
||||
from tinygrad.function import function
|
||||
from tinygrad import Tensor, GlobalCounters
|
||||
from tinygrad.uop.ops import UOp, Ops, KernelInfo
|
||||
|
||||
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_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.]])
|
||||
|
||||
@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)
|
||||
|
||||
@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.multi(0), device=devs)
|
||||
return Tensor.custom_kernel(c, a, fxn=double_kernel, grad_fxn=double_grad)[0]
|
||||
|
||||
a = Tensor.full((4, 4), 7.0).contiguous().shard(devs, axis=0)
|
||||
Tensor.realize(a)
|
||||
np.testing.assert_allclose(f(a).numpy(), 14.0)
|
||||
|
||||
def test_custom_kernel_precompile_further_compute(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"))
|
||||
|
||||
@function(precompile=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)[0]
|
||||
return c + 1
|
||||
|
||||
a = Tensor([1., 2., 3., 4.]).contiguous().realize()
|
||||
np.testing.assert_allclose(f(a).numpy(), [3., 5., 7., 9.])
|
||||
|
||||
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, 4739073)
|
||||
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
tinygrad_repo/test/unit/test_getitem_ops.py
Normal file
21
tinygrad_repo/test/unit/test_getitem_ops.py
Normal 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, 50_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, 500_000)
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
260
tinygrad_repo/test/unit/test_gguf.py
Normal file
260
tinygrad_repo/test/unit/test_gguf.py
Normal 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()
|
||||
183
tinygrad_repo/test/unit/test_gradient.py
Normal file
183
tinygrad_repo/test/unit/test_gradient.py
Normal file
@@ -0,0 +1,183 @@
|
||||
import unittest
|
||||
import numpy as np
|
||||
from tinygrad import Tensor
|
||||
from tinygrad.dtype import dtypes
|
||||
from tinygrad.uop.ops import UOp, KernelInfo
|
||||
|
||||
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_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
tinygrad_repo/test/unit/test_hashing.py
Normal file
109
tinygrad_repo/test/unit/test_hashing.py
Normal 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])
|
||||
|
||||
@unittest.skipUnless(Device.DEFAULT=="METAL", "slow")
|
||||
def test_sha3_224(self): self._test_preset("sha3_224", [143, 144])
|
||||
@unittest.skipUnless(Device.DEFAULT=="METAL", "slow")
|
||||
def test_sha3_256(self): self._test_preset("sha3_256", [135, 136])
|
||||
@unittest.skipUnless(Device.DEFAULT=="METAL", "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)
|
||||
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
tinygrad_repo/test/unit/test_hcq_graph.py
Normal file
44
tinygrad_repo/test/unit/test_hcq_graph.py
Normal 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(), UOp(Ops.DEVICE, arg=Device.DEFAULT))).call(UOp.new_buffer(Device.DEFAULT, 1, dtypes.float))
|
||||
cpu_call = UOp(Ops.PROGRAM, src=(UOp.sink(), UOp(Ops.DEVICE, arg="CPU"))).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()
|
||||
15
tinygrad_repo/test/unit/test_helpers.py
Normal file
15
tinygrad_repo/test/unit/test_helpers.py
Normal file
@@ -0,0 +1,15 @@
|
||||
import unittest
|
||||
import numpy as np
|
||||
from tinygrad.helpers import polyN, is_numpy_ndarray
|
||||
from tinygrad.tensor import Tensor
|
||||
|
||||
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()))
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
1393
tinygrad_repo/test/unit/test_indexing.py
Normal file
1393
tinygrad_repo/test/unit/test_indexing.py
Normal file
File diff suppressed because it is too large
Load Diff
126
tinygrad_repo/test/unit/test_invalid_tensor.py
Normal file
126
tinygrad_repo/test/unit/test_invalid_tensor.py
Normal file
@@ -0,0 +1,126 @@
|
||||
import unittest
|
||||
from tinygrad import Tensor
|
||||
from tinygrad.dtype import Invalid, dtypes
|
||||
from tinygrad.engine.realize import run_linear
|
||||
|
||||
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.copyin(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_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])
|
||||
|
||||
# tensor indexing uses reduce, so the entire result becomes invalid
|
||||
@unittest.expectedFailure
|
||||
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])
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
111
tinygrad_repo/test/unit/test_linalg.py
Normal file
111
tinygrad_repo/test/unit/test_linalg.py
Normal file
@@ -0,0 +1,111 @@
|
||||
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()
|
||||
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)
|
||||
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()
|
||||
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), Tensor.zeros(2, 2)):
|
||||
a = a.realize()
|
||||
U,S,V = a.svd()
|
||||
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)
|
||||
U,S,V = a.svd()
|
||||
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
tinygrad_repo/test/unit/test_llm_mla.py
Normal file
71
tinygrad_repo/test/unit/test_llm_mla.py
Normal 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
tinygrad_repo/test/unit/test_llm_moe.py
Normal file
100
tinygrad_repo/test/unit/test_llm_moe.py
Normal 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()
|
||||
163
tinygrad_repo/test/unit/test_llm_server.py
Normal file
163
tinygrad_repo/test/unit/test_llm_server.py
Normal file
@@ -0,0 +1,163 @@
|
||||
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
|
||||
|
||||
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)
|
||||
|
||||
class TestTransformerGenerate(unittest.TestCase):
|
||||
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):
|
||||
captured_inputs.append((tokens.shape, start_pos if isinstance(start_pos, int) else start_pos.val))
|
||||
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
|
||||
toks_shape = captured_inputs[0][0][-1]
|
||||
self.assertEqual(toks_shape.val if isinstance(toks_shape, UOp) else toks_shape, 4)
|
||||
self.assertEqual(captured_inputs[0][1], 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):
|
||||
captured_inputs.append((tokens.shape, start_pos if isinstance(start_pos, int) else start_pos.val))
|
||||
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
|
||||
toks_shape = captured_inputs[0][0][-1]
|
||||
self.assertEqual(toks_shape.val if isinstance(toks_shape, UOp) else toks_shape, 3)
|
||||
self.assertEqual(captured_inputs[0][1], 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):
|
||||
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):
|
||||
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
tinygrad_repo/test/unit/test_masked_tensor.py
Normal file
29
tinygrad_repo/test/unit/test_masked_tensor.py
Normal 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
tinygrad_repo/test/unit/test_metal_graph.py
Normal file
46
tinygrad_repo/test/unit/test_metal_graph.py
Normal 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(dtypes.weakint, 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()
|
||||
17
tinygrad_repo/test/unit/test_objc.py
Normal file
17
tinygrad_repo/test/unit/test_objc.py
Normal 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()
|
||||
194
tinygrad_repo/test/unit/test_randomness.py
Normal file
194
tinygrad_repo/test/unit/test_randomness.py
Normal file
@@ -0,0 +1,194 @@
|
||||
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_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)
|
||||
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(400)]
|
||||
torch_samples = [torch.tensor(w).multinomial(1, replacement=False).item() for _ in range(400)]
|
||||
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(400)])
|
||||
torch_draws = np.array([torch.tensor(w).multinomial(3, replacement=False).numpy() for _ in range(400)])
|
||||
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 = (128, 256, (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
tinygrad_repo/test/unit/test_realize_is_realize.py
Normal file
67
tinygrad_repo/test/unit/test_realize_is_realize.py
Normal 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 all(u.is_realized for u in t.uop.src)
|
||||
|
||||
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
tinygrad_repo/test/unit/test_rearrange_einops.py
Normal file
166
tinygrad_repo/test/unit/test_rearrange_einops.py
Normal 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
tinygrad_repo/test/unit/test_schedule_cache.py
Normal file
69
tinygrad_repo/test/unit/test_schedule_cache.py
Normal 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()
|
||||
151
tinygrad_repo/test/unit/test_setitem_schedule.py
Normal file
151
tinygrad_repo/test/unit/test_setitem_schedule.py
Normal file
@@ -0,0 +1,151 @@
|
||||
import unittest
|
||||
from tinygrad import Tensor, dtypes, GlobalCounters
|
||||
|
||||
class TestSetitemInto(unittest.TestCase):
|
||||
def test_setitem_into_unrealized(self):
|
||||
GlobalCounters.reset()
|
||||
t = Tensor.arange(4, dtype=dtypes.int32).reshape(2, 2)
|
||||
self.assertEqual(GlobalCounters.kernel_count, 0)
|
||||
t[1] = 5
|
||||
self.assertEqual(GlobalCounters.kernel_count, 0)
|
||||
t.realize()
|
||||
self.assertEqual(GlobalCounters.kernel_count, 1)
|
||||
self.assertEqual(GlobalCounters.global_mem, 16)
|
||||
t[1].realize()
|
||||
t.realize()
|
||||
self.assertEqual(GlobalCounters.kernel_count, 1)
|
||||
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]
|
||||
self.assertEqual(GlobalCounters.kernel_count, 0)
|
||||
w[1] = 99
|
||||
self.assertEqual(GlobalCounters.kernel_count, 0)
|
||||
w.realize()
|
||||
self.assertEqual(GlobalCounters.kernel_count, 1)
|
||||
self.assertEqual(GlobalCounters.global_mem, 4*4)
|
||||
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
|
||||
self.assertEqual(GlobalCounters.kernel_count, 0)
|
||||
t.realize()
|
||||
self.assertEqual(GlobalCounters.kernel_count, 1)
|
||||
self.assertEqual(GlobalCounters.global_mem, 4)
|
||||
t[1].realize()
|
||||
t.realize()
|
||||
self.assertEqual(GlobalCounters.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
|
||||
self.assertEqual(GlobalCounters.kernel_count, 0)
|
||||
t[1] = 5
|
||||
self.assertEqual(GlobalCounters.kernel_count, 0)
|
||||
t.realize()
|
||||
self.assertEqual(GlobalCounters.kernel_count, 1)
|
||||
self.assertLessEqual(GlobalCounters.global_mem, 32)
|
||||
t[1].realize()
|
||||
t.realize()
|
||||
self.assertEqual(GlobalCounters.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
|
||||
self.assertEqual(GlobalCounters.kernel_count, 0)
|
||||
t[1].realize()
|
||||
self.assertEqual(GlobalCounters.kernel_count, 1)
|
||||
self.assertEqual(GlobalCounters.global_mem, 4)
|
||||
t.realize()
|
||||
self.assertEqual(GlobalCounters.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
|
||||
self.assertEqual(GlobalCounters.kernel_count, 0)
|
||||
t[1].realize()
|
||||
self.assertEqual(GlobalCounters.kernel_count, 1)
|
||||
self.assertLessEqual(GlobalCounters.global_mem, 32)
|
||||
t[1].realize()
|
||||
t.realize()
|
||||
self.assertEqual(GlobalCounters.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
|
||||
self.assertEqual(GlobalCounters.kernel_count, 0)
|
||||
t.realize()
|
||||
self.assertEqual(GlobalCounters.kernel_count, 1)
|
||||
self.assertEqual(GlobalCounters.global_mem, 4*4)
|
||||
t[1].realize()
|
||||
t.realize()
|
||||
self.assertEqual(GlobalCounters.kernel_count, 1)
|
||||
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
|
||||
self.assertEqual(GlobalCounters.kernel_count, 0)
|
||||
t.realize()
|
||||
self.assertEqual(GlobalCounters.kernel_count, 1)
|
||||
self.assertEqual(GlobalCounters.global_mem, 4*4)
|
||||
t[1].realize()
|
||||
t.realize()
|
||||
self.assertEqual(GlobalCounters.kernel_count, 1)
|
||||
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
|
||||
self.assertEqual(GlobalCounters.kernel_count, 0)
|
||||
t.realize()
|
||||
self.assertEqual(GlobalCounters.kernel_count, 1)
|
||||
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
|
||||
self.assertEqual(GlobalCounters.kernel_count, 0)
|
||||
t.realize()
|
||||
self.assertEqual(GlobalCounters.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
|
||||
self.assertEqual(GlobalCounters.kernel_count, 0)
|
||||
t.realize()
|
||||
self.assertEqual(GlobalCounters.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
|
||||
self.assertEqual(GlobalCounters.kernel_count, 0)
|
||||
t.realize()
|
||||
self.assertEqual(GlobalCounters.kernel_count, 1)
|
||||
self.assertEqual(GlobalCounters.global_mem, 100*4) # full buffer written
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
40
tinygrad_repo/test/unit/test_shm_tensor.py
Normal file
40
tinygrad_repo/test/unit/test_shm_tensor.py
Normal 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
tinygrad_repo/test/unit/test_symbolic_tensor.py
Normal file
13
tinygrad_repo/test/unit/test_symbolic_tensor.py
Normal 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
tinygrad_repo/test/unit/test_system_pci_scan_bus.py
Normal file
28
tinygrad_repo/test/unit/test_system_pci_scan_bus.py
Normal 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
tinygrad_repo/test/unit/test_tar.py
Normal file
141
tinygrad_repo/test/unit/test_tar.py
Normal 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()
|
||||
91
tinygrad_repo/test/unit/test_tensor_data.py
Normal file
91
tinygrad_repo/test/unit/test_tensor_data.py
Normal file
@@ -0,0 +1,91 @@
|
||||
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(dtypes.int8, UOp.const(dtypes.float32, 1.0)).dtype, dtypes.int8)
|
||||
self.assertEqual(Tensor.const(dtypes.int32, UOp.variable("x", 1, 10).bind(5)).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_data_uop_device(self):
|
||||
uop = UOp.const(dtypes.float, 1.0, "DEVICE")
|
||||
self.assertEqual(Tensor(uop).device, "DEVICE")
|
||||
|
||||
if __name__ == '__main__':
|
||||
unittest.main()
|
||||
21
tinygrad_repo/test/unit/test_tensor_io.py
Normal file
21
tinygrad_repo/test/unit/test_tensor_io.py
Normal 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()
|
||||
81
tinygrad_repo/test/unit/test_tinyfs.py
Normal file
81
tinygrad_repo/test/unit/test_tinyfs.py
Normal file
@@ -0,0 +1,81 @@
|
||||
import json, math, os, socketserver, threading, unittest
|
||||
import numpy as np
|
||||
from tinygrad import Tensor, dtypes
|
||||
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) / Tensor.CHUNK_SIZE)):
|
||||
chunk = data[i*Tensor.CHUNK_SIZE:(i+1)*Tensor.CHUNK_SIZE].ljust(Tensor.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 = Tensor([1.0, 2.0, 3.0, 4.0]).fs_store().realize()
|
||||
self.assertEqual(h.shape, (16,))
|
||||
self.assertEqual(h.dtype, dtypes.uint8)
|
||||
|
||||
def test_store_deterministic(self):
|
||||
a = Tensor([1.0, 2.0, 3.0, 4.0]).fs_store().realize()
|
||||
b = Tensor([1.0, 2.0, 3.0, 4.0]).fs_store().realize()
|
||||
np.testing.assert_array_equal(a.numpy(), b.numpy())
|
||||
|
||||
def test_store_different_data(self):
|
||||
a = Tensor([1.0, 2.0, 3.0, 4.0]).fs_store().realize()
|
||||
b = Tensor([5.0, 6.0, 7.0, 8.0]).fs_store().realize()
|
||||
self.assertNotEqual(a.tolist(), b.tolist())
|
||||
|
||||
def test_roundtrip_uint8(self):
|
||||
arr = np.arange(256, dtype=np.uint8)
|
||||
loaded = Tensor(arr).fs_store().realize().fs_load(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=Tensor.CHUNK_SIZE + 1024, dtype=np.uint8)
|
||||
loaded = Tensor(arr).fs_store().realize().fs_load(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 = Tensor(arr).fs_store().realize()
|
||||
# the hash from fs_store should match the pure-Python hash_file reference
|
||||
padded = arr.tobytes().ljust(Tensor.CHUNK_SIZE, b'\0')
|
||||
self.assertEqual(h.data().tobytes(), hash_file(padded))
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
38
tinygrad_repo/test/unit/test_winograd.py
Normal file
38
tinygrad_repo/test/unit/test_winograd.py
Normal 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)
|
||||
Reference in New Issue
Block a user