From c5fa9eb36ef5b46e57cecb058482a7ec6c915e40 Mon Sep 17 00:00:00 2001 From: chenyu Date: Sat, 16 Dec 2023 01:25:44 -0500 Subject: [PATCH] int / List[int] data -> dtypes.int32 (#2789) --- test/test_dtype.py | 6 +++--- test/test_tensor.py | 8 ++++---- tinygrad/helpers.py | 4 ++-- tinygrad/tensor.py | 20 ++++++++------------ 4 files changed, 17 insertions(+), 21 deletions(-) diff --git a/test/test_dtype.py b/test/test_dtype.py index d4fe7c2cd6..5b64437f91 100644 --- a/test/test_dtype.py +++ b/test/test_dtype.py @@ -226,18 +226,18 @@ class TestHelpers(unittest.TestCase): class TestTypeSpec(unittest.TestCase): def test_creation(self): assert Tensor([]).dtype == Tensor.default_type - # assert Tensor([1]).dtype == dtypes.int + assert Tensor([1]).dtype == dtypes.int assert Tensor([1.1]).dtype == Tensor.default_type def test_const_full(self): assert Tensor.ones([2,3]).dtype == Tensor.default_type assert Tensor.zeros([2,3]).dtype == Tensor.default_type assert Tensor.full([2,3], 3.3).dtype == Tensor.default_type - # assert Tensor.full([2,3], 3).dtype == dtypes.int + assert Tensor.full([2,3], 3).dtype == dtypes.int def test_reduce_0d_default(self): assert Tensor.ones([2,3,0]).sum(2).dtype == Tensor.default_type - # assert Tensor.ones([2,3,0], dtype=dtypes.int).sum(2).dtype == dtypes.int + # assert Tensor.ones([2,3,0], dtype=dtypes.int).sum(2).dtype == dtypes.int # requires reduceop acc fix def test_arange(self): assert Tensor.arange(5).dtype == dtypes.int32 diff --git a/test/test_tensor.py b/test/test_tensor.py index 6e8b4da611..f97e8441c0 100644 --- a/test/test_tensor.py +++ b/test/test_tensor.py @@ -152,7 +152,7 @@ class TestTinygrad(unittest.TestCase): except: raise finally: Tensor.rand = original_rand - def test_zeros_like_has_same_dtype(self): + def test_zeros_like_has_same_dtype_and_shape(self): for datatype in [dtypes.float16, dtypes.float32, dtypes.int8, dtypes.int32, dtypes.int64, dtypes.uint8]: a = Tensor([1, 2, 3], dtype=datatype) b = Tensor.zeros_like(a) @@ -161,7 +161,7 @@ class TestTinygrad(unittest.TestCase): a = Tensor([1, 2, 3]) b = Tensor.zeros_like(a, dtype=dtypes.int8) - assert a.dtype != b.dtype and a.dtype == dtypes.float32 and b.dtype == dtypes.int8, "a.dtype should be float and b.dtype should be char" + assert a.dtype == dtypes.int32 and b.dtype == dtypes.int8, "a.dtype should be int and b.dtype should be char" assert a.shape == b.shape, f"shape mismatch (Tensor.zeros_like){a.shape} != (torch){b.shape}" def test_ones_like_has_same_dtype_and_shape(self): @@ -173,7 +173,7 @@ class TestTinygrad(unittest.TestCase): a = Tensor([1, 2, 3]) b = Tensor.ones_like(a, dtype=dtypes.int8) - assert a.dtype != b.dtype and a.dtype == dtypes.float32 and b.dtype == dtypes.int8, "a.dtype should be float and b.dtype should be char" + assert a.dtype == dtypes.int32 and b.dtype == dtypes.int8, "a.dtype should be int and b.dtype should be char" assert a.shape == b.shape, f"shape mismatch (Tensor.ones_like){a.shape} != (torch){b.shape}" def test_ndim(self): @@ -240,7 +240,7 @@ class TestTinygrad(unittest.TestCase): def test_tensor_list_dtype(self): arr = [1] - assert Tensor(arr).dtype == Tensor.default_type + assert Tensor(arr).dtype == dtypes.int32 assert Tensor(arr, dtype=dtypes.float32).dtype == dtypes.float32 assert Tensor(arr, dtype=dtypes.float64).dtype == dtypes.float64 diff --git a/tinygrad/helpers.py b/tinygrad/helpers.py index 62d124e2de..9f6c65fc64 100644 --- a/tinygrad/helpers.py +++ b/tinygrad/helpers.py @@ -3,7 +3,7 @@ import os, functools, platform, time, re, contextlib, operator, hashlib, pickle, import numpy as np from urllib import request from tqdm import tqdm -from typing import Dict, Tuple, Union, List, NamedTuple, Final, ClassVar, Optional, Iterable, Any, TypeVar, TYPE_CHECKING, Callable, Set +from typing import Dict, Tuple, Union, List, NamedTuple, Final, ClassVar, Optional, Iterable, Any, TypeVar, TYPE_CHECKING, Callable, Set, Sequence if TYPE_CHECKING: # TODO: remove this and import TypeGuard from typing once minimum python supported version is 3.10 from typing_extensions import TypeGuard @@ -20,7 +20,7 @@ def dedup(x:Iterable[T]): return list(dict.fromkeys(x)) # retains list order def argfix(*x): return tuple(x[0]) if x and x[0].__class__ in (tuple, list) else x def argsort(x): return type(x)(sorted(range(len(x)), key=x.__getitem__)) # https://stackoverflow.com/questions/3382352/equivalent-of-numpy-argsort-in-basic-python def all_same(items:List[T]): return all(x == items[0] for x in items) -def all_int(t: Tuple[Any, ...]) -> TypeGuard[Tuple[int, ...]]: return all(isinstance(s, int) for s in t) +def all_int(t: Sequence[Any]) -> TypeGuard[Tuple[int, ...]]: return all(isinstance(s, int) for s in t) def colored(st, color:Optional[str], background=False): return f"\u001b[{10*background+60*(color.upper() == color)+30+['black', 'red', 'green', 'yellow', 'blue', 'magenta', 'cyan', 'white'].index(color.lower())}m{st}\u001b[0m" if color is not None else st # replace the termcolor library with one line # noqa: E501 def ansistrip(s:str): return re.sub('\x1b\\[(K|.*?m)', '', s) def ansilen(s:str): return len(ansistrip(s)) diff --git a/tinygrad/tensor.py b/tinygrad/tensor.py index 5324081886..50b98ad2f9 100644 --- a/tinygrad/tensor.py +++ b/tinygrad/tensor.py @@ -59,19 +59,15 @@ class Tensor: # internal variables used for autograd graph construction self._ctx: Optional[Function] = None if isinstance(data, LazyBuffer): assert dtype is None or dtype == data.dtype, "dtype doesn't match, and casting isn't supported" - elif isinstance(data, (int, float)): - data = LazyBuffer.loadop(LoadOps.CONST, tuple(), dtype or Tensor.default_type, device, data) - elif data is None or data.__class__ is list: - assert dtype is None or dtype.np is not None, f"{dtype} doesn't have a numpy dtype" - data = LazyBuffer.fromCPU(np.array([] if data is None else data, dtype=(dtype or Tensor.default_type).np)) - elif isinstance(data, bytes): - data = LazyBuffer.fromCPU(np.frombuffer(data, np.uint8)) + elif isinstance(data, int): data = LazyBuffer.loadop(LoadOps.CONST, tuple(), dtype or dtypes.int32, device, data) + elif isinstance(data, float): data = LazyBuffer.loadop(LoadOps.CONST, tuple(), dtype or Tensor.default_type, device, data) + elif isinstance(data, bytes): data = LazyBuffer.fromCPU(np.frombuffer(data, np.uint8)) + elif data is None: data = LazyBuffer.fromCPU(np.array([], dtype=(dtype or Tensor.default_type).np)) + elif isinstance(data, list): + data = LazyBuffer.fromCPU(np.array(data, (dtype or (dtypes.int32 if data and all_int(data) else Tensor.default_type)).np)) elif isinstance(data, np.ndarray): - assert dtype is None or dtype.np is not None, f"{dtype} doesn't have a numpy dtype" - if data.shape == (): - data = LazyBuffer.loadop(LoadOps.CONST, tuple(), dtype or dtypes.from_np(data.dtype), device, data.item()) - else: - data = LazyBuffer.fromCPU(data.astype(dtype.np) if dtype is not None and dtype.np is not None else data) + if data.shape == (): data = LazyBuffer.loadop(LoadOps.CONST, tuple(), dtype or dtypes.from_np(data.dtype), device, data.item()) + else: data = LazyBuffer.fromCPU(data.astype(dtype.np) if dtype is not None and dtype.np is not None else data) # data is a LazyBuffer, but it might be on the wrong device if not isinstance(data, LazyBuffer): raise RuntimeError(f"can't create Tensor from {data!r} with type {type(data)}")