diff --git a/tinygrad/llm/gguf.py b/tinygrad/llm/gguf.py index 39e56ade38..de1482515e 100644 --- a/tinygrad/llm/gguf.py +++ b/tinygrad/llm/gguf.py @@ -38,7 +38,7 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor: def q_to_uint8(t: Tensor, b: int) -> Tensor: # TODO: rewrite with arange? shift_tensor, bitmask = Tensor.const(t.dtype, tuple(2**(i*b) for i in range(8//b))), 0xff >> (8 - b) - return t.unsqueeze(-1).expand((*t.shape,8//b)).div(shift_tensor, rounding_mode="trunc").bitwise_and(bitmask).transpose(-1, -2).flatten(-2) + return t.unsqueeze(-1).div(shift_tensor, rounding_mode="trunc").bitwise_and(bitmask).transpose(-1, -2).flatten(-2) if (nelements_nbytes := _GGML_QUANT.get(ggml_type)) is not None: from tinygrad.runtime.autogen import ggml_common as _ggml @@ -68,7 +68,7 @@ def ggml_data_to_tensor(t: Tensor, n: int, ggml_type: int) -> Tensor: if ggml_type == 14: xl, xh = q_to_uint8(blocks[:,:128].reshape((-1, 2, 64)), 4), q_to_uint8(blocks[:,128:192].reshape((-1, 2, 32)), 2).lshift(4) scales = blocks[:,192:208].bitcast(dtypes.int8).unsqueeze(-1).expand((-1, 16, 16)).reshape((-1, 256)) - d = blocks[:,-2:].bitcast(dtypes.float16).cast(dtypes.float32).expand((-1, 256)) + d = blocks[:,-2:].bitcast(dtypes.float16).cast(dtypes.float32) return d * (xl.bitwise_or(xh).bitcast(dtypes.int8) - 32).flatten(-2) * scales if ggml_type == 18: d = blocks[:, :2].bitcast(dtypes.float16).cast(dtypes.float32).reshape((-1, 1, 1, 1)) diff --git a/tinygrad/mixin/op.py b/tinygrad/mixin/op.py index f523215d99..7474debbc5 100644 --- a/tinygrad/mixin/op.py +++ b/tinygrad/mixin/op.py @@ -440,7 +440,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin): sz = merge_dicts([dict(zip(s, x.shape)) for s, x in zip(inputs, xs)]) alpha = sorted(sz) # align all tensors to alphabet, multiply, sum non-output, permute to output order - xs = [x.permute(*[s.index(c) for c in sorted(s)]).reshape([sz[c] if c in s else 1 for c in alpha]).expand([sz[c] for c in alpha]) if s else x + xs = [x.permute(*[s.index(c) for c in sorted(s)]).reshape([sz[c] if c in s else 1 for c in alpha]) if s else x for s, x in zip(inputs, xs)] return xs[0].uprod(*xs[1:]).sum([i for i,c in enumerate(alpha) if c not in rhs], dtype=dtype).permute(argsort(argsort(list(rhs)))) @@ -831,7 +831,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin): if self.ndim == 0: return self x = self.transpose(axis, -1) last_dim_size = x.shape[-1] - x_unsqueezed = x.unsqueeze(-2).expand((None,)*(self.ndim-1)+(last_dim_size, None)) + x_unsqueezed = x.unsqueeze(-2) x_cummax = x.cummax(-1)[0].detach() mask = type(self).ones(last_dim_size, last_dim_size, buffer=False, dtype=dtypes.bool).tril() ret = mask.where(x_unsqueezed - x_cummax.unsqueeze(-1), self.dtype.min).exp().sum(-1).log() + x_cummax diff --git a/tinygrad/mixin/rand.py b/tinygrad/mixin/rand.py index b621633591..003af96cd7 100644 --- a/tinygrad/mixin/rand.py +++ b/tinygrad/mixin/rand.py @@ -269,7 +269,7 @@ class RandMixin(OpMixin): if replacement or num_samples == 1: cdf = (cw := weight.cumsum(1).float()) / cw[:, -1].unsqueeze(1) unif_samples = type(self).rand(num_samples, cdf.shape[0], 1).to(self.device) # type: ignore[attr-defined] - indices = (unif_samples.expand((-1, -1, cdf.shape[1])) >= cdf).sum(2).permute((1, 0)) + indices = (unif_samples >= cdf).sum(2).permute((1, 0)) else: # Efraimidis-Spirakis indices = (weight.rand_like(dtype=dtypes.float32).log2() / weight).topk(num_samples, dim=1)[1] diff --git a/tinygrad/nn/__init__.py b/tinygrad/nn/__init__.py index c4732bcd0d..203a73fcca 100644 --- a/tinygrad/nn/__init__.py +++ b/tinygrad/nn/__init__.py @@ -40,7 +40,7 @@ class BatchNorm: def calc_stats(self, x:Tensor) -> tuple[Tensor, Tensor]: shape_mask: list[int] = [1, -1, *([1]*(x.ndim-2))] - if self.track_running_stats and not TRAINING: return self.running_mean, self.running_var.reshape(shape=shape_mask).expand(x.shape) + if self.track_running_stats and not TRAINING: return self.running_mean, self.running_var.reshape(shape=shape_mask) # This requires two full memory accesses to x # https://github.com/pytorch/pytorch/blob/c618dc13d2aa23625cb0d7ada694137532a4fa33/aten/src/ATen/native/cuda/Normalization.cuh # There's "online" algorithms that fix this, like https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Welford's_Online_algorithm