From 590268fa03bf4bb6bac8651fcbe9538eaf95d3fe Mon Sep 17 00:00:00 2001 From: chenyu Date: Wed, 3 Jan 2024 11:45:28 -0500 Subject: [PATCH] out_tokens -> grouped in linearizer (#2989) no more token now --- tinygrad/codegen/linearizer.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/tinygrad/codegen/linearizer.py b/tinygrad/codegen/linearizer.py index 08d1c8cc01..c03ae1c404 100644 --- a/tinygrad/codegen/linearizer.py +++ b/tinygrad/codegen/linearizer.py @@ -158,11 +158,11 @@ class Linearizer(Kernel): _idx = k[:upcast_dim[0]] + (float4_expand[0],) + k[upcast_dim[0]+1:] grouped_store_offset[_idx].append(store_offset[k]) store_offset_new = {} - for k,out_tokens in grouped_store_offset.items(): - amt = len(out_tokens) + for k,grouped in grouped_store_offset.items(): + amt = len(grouped) idx, valid = self.sts[i].expr_idxs(k) assert idx == ((idx//amt)*amt), "float4 stores are always aligned" - store_offset_new[k] = self.uop(UOps.CAST, buf.dtype.vec(amt), tuple(out_tokens)) + store_offset_new[k] = self.uop(UOps.CAST, buf.dtype.vec(amt), tuple(grouped)) store_offset = store_offset_new stores = []