fix assign

test/external/external_test_simple_tokenizer.py

@@ -1,50 +1,41 @@
-import functools, multiprocessing
 from transformers import AutoTokenizer
 from datasets import load_dataset
-from tinygrad.apps.llm import SimpleTokenizer
+from tinygrad.apps.llm import SimpleTokenizer, gpt2_decode_vocab, get_llama_re
 from tinygrad.helpers import tqdm, getenv, partition
 
-@functools.cache
-def get_tokenizers():
-  print("getting tokenizers")
-  base_tokenizer = AutoTokenizer.from_pretrained("NousResearch/Meta-Llama-3-8B-Instruct")
-  special_tokens, normal_tokens = partition(((t, tid) for t, tid in base_tokenizer.vocab.items()), lambda e: e[1] in base_tokenizer.all_special_ids)
-  simple_tokenizer = SimpleTokenizer(dict(normal_tokens), dict(special_tokens))
-  return base_tokenizer, simple_tokenizer
-
-def test_tokenize(samp) -> bool:
-  base_tokenizer, simple_tokenizer = get_tokenizers()
-  idx, txt = samp
-  try: simple_tokens = tuple(simple_tokenizer.encode(txt))
-  except RuntimeError: simple_tokens = ()
-  base_tokens = tuple(base_tokenizer.encode(txt, add_special_tokens=False))
-  if simple_tokens != base_tokens:
-    print(f"tokens mismatch at index: {idx}.\n")
-    color_codes = [91, 92, 94, 93, 95]
-    def color_tokens(tids):
-      return "".join(f"\033[{color_codes[i%len(color_codes)]}m{base_tokenizer.decode([t])}" for i, t in enumerate(tids)) + "\033[0m"
-    print("simple:  ", color_tokens(simple_tokens))
-    print("official:", color_tokens(base_tokens) + "\n")
-    return False
-  if simple_tokenizer.decode(simple_tokens) != txt:
-    print(f"decode mismatch at {idx}")
-    return False
-  return True
-
 # use ALLOW_FAILED=-1 to go over the entire dataset without printing.
 if __name__ == "__main__":
-  print("loading datasets")
-  ds = load_dataset("OpenAssistant/oasst1")
-  loaded_ds = [(idx, el["text"]) for idx, el in enumerate(ds["train"])]
-  print(f"loaded {len(loaded_ds)}")
+  base_tokenizer = AutoTokenizer.from_pretrained("NousResearch/Meta-Llama-3-8B-Instruct")
+  special_tokens, normal_tokens = partition(((t, tid) for t, tid in base_tokenizer.vocab.items()),
+    lambda e: e[1] in base_tokenizer.all_special_ids)
+  inv_vocab = { tid: word for word, tid in base_tokenizer.get_vocab().items() }
+  simple_tokenizer = SimpleTokenizer(get_llama_re(), gpt2_decode_vocab(dict(normal_tokens)), dict(special_tokens))
 
+  color_codes = [ 91, 92, 94, 93, 95 ]
+  def color_tokens(tids):
+    return "".join(f"\033[{color_codes[i%len(color_codes)]}m{base_tokenizer.decode([t])}" for i, t in enumerate(tids)) + "\033[0m"
+
+  ds = load_dataset("OpenAssistant/oasst1")
   allow_failed = getenv("ALLOW_FAILED", 10)
+
   fail_count, total = 0, 0
-  with multiprocessing.Pool(16) as pool:
-    for good in tqdm(pool.imap_unordered(test_tokenize, loaded_ds), total=len(loaded_ds)):
-      total += 1
-      if not good:
-        fail_count += 1
-        allow_failed -= 1
-        if allow_failed == 0: break
-    print(f"{fail_count}/{total} samples are inconsistent with the official tokenizer.")
+
+  for idx, el in enumerate(tqdm(ds["train"])):
+    total += 1
+
+    try: simple_tokens = tuple(simple_tokenizer.encode(el["text"]))
+    except RuntimeError: simple_tokens = ()
+    base_tokens = tuple(base_tokenizer.encode(el["text"], add_special_tokens=False))
+
+    if simple_tokens != base_tokens:
+      fail_count += 1
+      allow_failed -= 1
+
+      if allow_failed >= 0:
+        print(f"tokens mismatch at index: {idx}.\n")
+
+        print("simple:  ", color_tokens(simple_tokens))
+        print("official:", color_tokens(base_tokens) + "\n")
+
+      if allow_failed == 0: break
+  print(f"{fail_count}/{total} samples are inconsistent with the official tokenizer.")

test/test_outerworld.py

@@ -1,63 +1,7 @@
 import unittest
-from tinygrad import Tensor, UOp, Variable, nn
+from tinygrad import Tensor, UOp
 from tinygrad.uop.ops import AxisType, Ops
 
-class TestOuterworldTrain(unittest.TestCase):
-  @Tensor.train()
-  def test_train(self):
-    # same example over and over
-    X = Tensor.rand(1, 32).expand(16,32).contiguous()
-    Y = Tensor.rand(1, 1).expand(16,1).contiguous()
-
-    layer = nn.Linear(32, 1, bias=False)
-    opt = nn.optim.SGD(nn.state.get_parameters(layer))
-    Tensor.realize(X, Y, *nn.state.get_parameters(layer))
-
-    print("train")
-
-    # if everything is correct, this should be a 16 step training loop
-    steps = UOp.range(16, -1)
-    opt.zero_grad()
-    loss = (layer(X[steps]) - Y[steps]).square().mean().backward()
-    sched = opt.schedule_step() # TODO: does this need to know anything about steps?
-    # NOTE: this can't work. the inputs to layer are not the assign, need to run twice for the fixed point?
-    all_losses = Tensor.realize(loss.reshape(1).expand(steps).contiguous(), *sched)
-    print(all_losses.numpy())
-
-#@unittest.skip("TODO: understand assign")
-class TestOuterworldAssign(unittest.TestCase):
-  def test_triple_add_inner(self):
-    t = Tensor.zeros(5).contiguous().realize()
-    t2 = Tensor.ones(3).contiguous().realize()
-    a = UOp.range(3, -1)
-    t = t.reshape(1,5).expand(a+1,5)[a].assign(t+t2[a])
-    self.assertListEqual(t.tolist(), [3,3,3,3,3])
-
-  def test_triple_add_outer(self):
-    t = Tensor.zeros(5).contiguous().realize()
-    t2 = Tensor.ones(3).contiguous().realize()
-
-    # OUTER is a loop at the schedule level
-    a = UOp.range(3, -1, AxisType.OUTER)
-    va = Variable("loop", 0, 2).bind(a)
-    t = t.assign(t+t2[va])
-    t = Tensor(UOp(Ops.ENDRANGE, dtype=t.uop.dtype, src=(a, t.uop)))
-
-    self.assertListEqual(t.tolist(), [3,3,3,3,3])
-
-  def test_triple_gemm(self):
-    x = Tensor.rand(1, 16).realize()
-    W = Tensor.rand(3, 16, 16).realize()
-
-    #manual = (x @ W[0] @ W[1] @ W[2]).contiguous().realize()
-
-    a = UOp.range(3, -1)
-
-    out = (x @ W[a]).contiguous()
-    t = Tensor(UOp(Ops.ASSIGN, dtype=out.uop.dtype, src=(x.uop, out.uop, a)))
-    #t = Tensor(UOp(Ops.REDUCE, dtype=out.uop.dtype, src=(out.uop, x.uop, a), arg=Ops.NOOP))
-    t.realize()
-
 class TestOuterworldReduce(unittest.TestCase):
   def test_reduce(self):
     x = Tensor.ones(5, 5).contiguous()
@@ -96,7 +40,6 @@ class TestOuterworld(unittest.TestCase):
     # passthrough ranges
     a = UOp.range(10, -1)
     sel = t[9-a]
-    assert sel.shape == (10,)
     cpy = sel.reshape(1, 10).expand(a, 10).contiguous().realize()
 
     self.assertTrue((t.flip(0)==cpy).all().item())

test/unit/test_llm_tokenizer.py

@@ -1,21 +1,19 @@
 import unittest, base64, functools, sys
-from tinygrad.apps.llm import SimpleTokenizer
+from tinygrad.apps.llm import SimpleTokenizer, get_llama_re
 from tinygrad.helpers import fetch
 
 @unittest.skipIf(sys.platform == 'win32', "fetch race condition on Windows")
 class TestLLMTokenizer(unittest.TestCase):
+  @functools.cached_property
+  def basic_tok(self): return SimpleTokenizer(".*", { b"a": 0, b"b": 1, b"c": 2, b"ab": 3, b"bc": 4 }, { "<x>": 5, "<y>": 6, "<z>": 7 })
+
   @functools.cached_property
   def llama_tok(self):
     # from https://github.com/tinygrad/tinygrad/blob/e0106b6b257ebc003eb3694144e3e198f7d8cc37/examples/llama3.py#L14
     model_file = fetch("https://huggingface.co/bofenghuang/Meta-Llama-3-8B/resolve/main/original/tokenizer.model")
     with open(model_file, "rt") as fd:
-      str_vocab = [line.split(maxsplit=1) for line in fd.read().splitlines() if line]
-
-      # https://github.com/openai/gpt-2/blob/9b63575ef42771a015060c964af2c3da4cf7c8ab/src/encoder.py#L9
-      bs = [*range(33, 127), *range(161, 173), *range(174, 256)]  # bytes that map to themselves
-      _byte_decoder = {chr(b): b for b in bs} | {chr(256+i): b for i,b in enumerate(b for b in range(256) if b not in bs)}
-      _byte_encoder = {v:k for k,v in _byte_decoder.items()}
-      normal_tokens = {''.join([_byte_encoder[x] for x in base64.b64decode(stok)]): int(srank) for stok, srank in str_vocab}
+      str_vocab = [ line.split(maxsplit=1) for line in fd.read().splitlines() if line ]
+      normal_tokens = { base64.b64decode(stok): int(srank) for stok, srank in str_vocab }
 
     special_tokens = [
       "<|begin_of_text|>",
@@ -29,12 +27,22 @@ class TestLLMTokenizer(unittest.TestCase):
       "<|reserved_special_token_4|>",
       "<|eot_id|>",
     ] + [ f"<|reserved_special_token_{i}|>" for i in range(5, 256 - 5) ]
-    return SimpleTokenizer(normal_tokens, {token: len(normal_tokens) + i for i, token in enumerate(special_tokens)})
+    return SimpleTokenizer(get_llama_re(), normal_tokens, { token: len(normal_tokens) + i for i, token in enumerate(special_tokens) })
 
   def _test_coding(self, tok: SimpleTokenizer, text: str, expected_tokens: list[int]):
     self.assertEqual(tok.encode(text), expected_tokens)
     self.assertEqual(tok.decode(expected_tokens), text)
 
+  def test_abc(self): self._test_coding(self.basic_tok, "abc", [ 3, 2 ])
+  def test_abbc(self): self._test_coding(self.basic_tok, "abbc", [ 3, 4 ])
+  def test_aabbbcc(self): self._test_coding(self.basic_tok, "aabbbcc", [ 0, 3, 1, 4, 2 ])
+  def test_specials1(self): self._test_coding(self.basic_tok, "a<x>a<y>a<z>a", [ 0, 5, 0, 6, 0, 7, 0 ])
+  def test_specials2(self): self._test_coding(self.basic_tok, "<x>a<y>a<z>", [ 5, 0, 6, 0, 7 ])
+  def test_invalid_token(self):
+    with self.assertRaises(RuntimeError): self._test_coding(self.basic_tok, "L", [])
+
+  def test_no_specials(self): self._test_coding(SimpleTokenizer(".*", { bytes([i]): i for i in range(256) }, {}), "abc", [97, 98, 99])
+
   # NOTE: the correct tokenization for this can only be found by looking up the text chunk in the vocab, not by applying merges
   def test_llama_early_tokenize(self): self._test_coding(self.llama_tok, " например", [ 111797 ])
 

tinygrad/apps/llm.py

@@ -1,55 +1,58 @@
 from __future__ import annotations
-import sys, argparse, typing, re, unicodedata
+import sys, argparse, typing, re, itertools, unicodedata
 from tinygrad import Tensor, nn, UOp, TinyJit, getenv, helpers
 
+def gpt2_decode_vocab(voc: dict[str, int]): # https://github.com/openai/gpt-2/blob/9b63575ef42771a015060c964af2c3da4cf7c8ab/src/encoder.py#L9
+  c2b = { chr(cp): cp for cp in itertools.chain(range(ord("!"), ord("~")+1), range(ord("¡"), ord("¬")+1), range(ord("®"), ord("ÿ")+1)) }
+  c2b.update({ chr(256+off): cp for off, cp in enumerate(cp for cp in range(256) if chr(cp) not in c2b) })
+  return { bytes(c2b[c] for c in tok): tid for tok, tid in voc.items() }
+
+def get_llama_re():
+  def ucat_range(pre: str): return "".join(re.escape(chr(cp)) for cp in range(sys.maxunicode + 1) if unicodedata.category(chr(cp)).startswith(pre))
+  r_ws, r_p_N, r_p_L = r"\t\n\x0b\x0c\r\x85" + ucat_range("Z"), ucat_range("N"), ucat_range("L")
+  # https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L286
+  return "(?i:'s|'t|'re|'ve|'m|'ll|'d)|" + \
+    f"[^\\r\\n{r_p_N}{r_p_L}]?[{r_p_L}]+|[{r_p_N}]{{1,3}}| ?[^{r_ws}{r_p_N}{r_p_L}]+[\\r\\n]*|[{r_ws}]*[\\r\\n]+|[{r_ws}]+(?![^{r_ws}])|[{r_ws}]+"
+
 class SimpleTokenizer:
-  def __init__(self, normal_tokens:dict[str, int], special_tokens:dict[str, int]):
-    # https://github.com/openai/gpt-2/blob/9b63575ef42771a015060c964af2c3da4cf7c8ab/src/encoder.py#L9
-    bs = [*range(33, 127), *range(161, 173), *range(174, 256)]  # bytes that map to themselves
-    self._byte_decoder = {chr(b): b for b in bs} | {chr(256+i): b for i,b in enumerate(b for b in range(256) if b not in bs)}
-
-    # https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L286
-    def ucat_range(pre: str): return "".join(re.escape(chr(cp)) for cp in range(sys.maxunicode + 1) if unicodedata.category(chr(cp)).startswith(pre))
-    r_ws, r_p_N, r_p_L = r"\t\n\x0b\x0c\r\x85" + ucat_range("Z"), ucat_range("N"), ucat_range("L")
-    self._split_to_word = re.compile("(?i:'s|'t|'re|'ve|'m|'ll|'d)|" + \
-      f"[^\\r\\n{r_p_N}{r_p_L}]?[{r_p_L}]+|[{r_p_N}]{{1,3}}| ?[^{r_ws}{r_p_N}{r_p_L}]+[\\r\\n]*|[{r_ws}]*[\\r\\n]+|[{r_ws}]+(?![^{r_ws}])|[{r_ws}]+")
-    self._split_to_sentence = re.compile("|".join(re.escape(tok) for tok in special_tokens.keys()) if special_tokens else r"(?!)")
-
-    self._normal_tokens = {bytes(self._byte_decoder[c] for c in tok): tid for tok, tid in normal_tokens.items()}
-    self._special_tokens = special_tokens
-    self._tok2bytes = {tid: tok for tok, tid in self._normal_tokens.items()} | {tid: tok.encode() for tok, tid in self._special_tokens.items()}
+  def __init__(self, pat: str, normal_tokens: dict[bytes, int], special_tokens: dict[str, int]):
+    self._normal_tokens, self._special_tokens, self._pat = normal_tokens, special_tokens, re.compile(pat)
+    self._tok2str = { tid: tok.encode() for tok, tid in special_tokens.items() } | { tid: tok for tok, tid in normal_tokens.items()  }
+    self._special_re = re.compile("|".join(re.escape(tok) for tok in self._special_tokens.keys()) if special_tokens else r"(?!)")
 
   @staticmethod
-  def from_gguf_kv(kv:dict):
+  def from_gguf_kv(kv: dict):
     # https://github.com/ggml-org/llama.cpp/blob/94933c8c2eeaa9a7983e3f6c08af76bd86724094/src/llama-vocab.cpp#L1818-L1820
     if kv["tokenizer.ggml.pre"] not in ("llama3","llama-v3","llama-bpe"): raise ValueError(f"Invalid tokenizer preset '{kv['tokenizer.ggml.pre']}'")
     vocab: typing.Iterable[tuple[str, int]] = ((tok, idx) for idx, tok in enumerate(kv["tokenizer.ggml.tokens"]))
     normal_tokens, special_tokens = helpers.partition(vocab, lambda e: kv["tokenizer.ggml.token_type"][e[1]] == 1)
-    return SimpleTokenizer(dict(normal_tokens), dict(special_tokens))
+    return SimpleTokenizer(get_llama_re(), gpt2_decode_vocab(dict(normal_tokens)), dict(special_tokens))
 
-  def _encode_word(self, word:bytes) -> list[int]:
-    if (early_token:=self._normal_tokens.get(word)) is not None: return [early_token]
-    parts = [bytes([b]) for b in word]
-    # greedily merge any parts that we can
-    while True:
-      i = min([(sys.maxsize, -1)] + [(self._normal_tokens.get(parts[j]+parts[j+1], sys.maxsize), j) for j in range(len(parts)-1)])[1]
-      if i == -1: break
-      parts[i:i+2] = [parts[i] + parts[i+1]]
-    try: return [self._normal_tokens[p] for p in parts]
-    except KeyError: raise RuntimeError("token not found")
-  def _encode_sentence(self, chunk:str) -> list[int]:
-    return [tok for word in self._split_to_word.findall(chunk) for tok in self._encode_word(word.encode())]
-  def encode(self, text:str) -> list[int]:
+  def encode(self, text: str):
     tokens: list[int] = []
     pos = 0
-    for match in self._split_to_sentence.finditer(text):
+    for match in self._special_re.finditer(text):
       tokens.extend(self._encode_sentence(text[pos:match.start(0)]) + [self._special_tokens[text[match.start(0):match.end(0)]]])
       pos = match.end(0)
     return tokens + self._encode_sentence(text[pos:])
 
-  def decode(self, ids:list[int]) -> str: return b''.join(self._tok2bytes[tid] for tid in ids).decode()
+  def decode(self, ids: list[int]) -> str: return b''.join(self._tok2str[tid] for tid in ids).decode()
   def role(self, role:str): return self.encode("<|start_header_id|>" + role + "<|end_header_id|>\n\n")
 
+  def _encode_sentence(self, chunk: str): return [ tok for word in self._pat.findall(chunk) for tok in self._encode_word(word.encode()) ]
+  def _encode_word(self, word: bytes):
+    if (early_token:=self._normal_tokens.get(word)) is not None: return [early_token]
+    parts = [word[i:i+1] for i in range(len(word))]
+    while True:
+      min_tid, min_idx = 2**32, -1
+      for idx, (p1, p2) in enumerate(zip(parts[:-1], parts[1:])):
+        tid = self._normal_tokens.get(p1 + p2, min_tid)
+        if tid < min_tid: min_tid, min_idx = tid, idx
+      if min_idx == -1: break
+      parts = parts[:min_idx] + [parts[min_idx] + parts[min_idx+1]] + parts[min_idx+2:]
+    try: return [ self._normal_tokens[p] for p in parts ]
+    except KeyError: raise RuntimeError("token not found")
+
 def apply_rope(x:Tensor, start_pos:int|UOp, base:float = 10000.0) -> Tensor:
   B, H, T, Hd = x.shape
   assert isinstance(Hd, int) and (Hd & 1) == 0, "RoPE requires an even head dimension"

tinygrad/codegen/opt/__init__.py

@@ -18,8 +18,8 @@ class Opt:
 
 axis_letters = {AxisType.GLOBAL: "g", AxisType.THREAD: "t", AxisType.LOCAL: "l", AxisType.WARP: "w", AxisType.LOOP: "L", AxisType.UPCAST: "u",
                 AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r"}
-axis_colors = {AxisType.OUTER: "GREEN", AxisType.GLOBAL: "blue", AxisType.THREAD: "BLUE", AxisType.LOCAL: "cyan", AxisType.WARP: "CYAN",
-               AxisType.LOOP: "WHITE", AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "RED", AxisType.REDUCE: "red", AxisType.UNROLL: "magenta"}
+axis_colors = {AxisType.GLOBAL: "blue", AxisType.THREAD: "BLUE", AxisType.LOCAL: "cyan", AxisType.WARP: "CYAN", AxisType.LOOP: "WHITE",
+               AxisType.UPCAST: "yellow", AxisType.GROUP_REDUCE: "RED", AxisType.REDUCE: "red", AxisType.UNROLL: "magenta"}
 
 class KernelOptError(Exception): pass
 def check(cond:bool, msg:str=""):

tinygrad/codegen/opt/postrange.py

@@ -13,8 +13,8 @@ from tinygrad.renderer import Renderer
 remove_tags = PatternMatcher([(UPat(GroupOp.All, name="x"), lambda x: x.replace(tag=None) if x.tag is not None else None)])
 
 # NOTE: LOCAL and GROUP_REDUCE have the same priority. the order here matters
-axis_to_pos = {AxisType.OUTER: -2, AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2,
-               AxisType.UPCAST: 3, AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}
+axis_to_pos = {AxisType.LOOP: -1, AxisType.THREAD: 0, AxisType.GLOBAL: 0, AxisType.WARP: 1, AxisType.LOCAL: 2, AxisType.UPCAST: 3,
+               AxisType.GROUP_REDUCE: 2, AxisType.REDUCE: 4, AxisType.UNROLL: 5}
 
 class Scheduler:
   def __init__(self, ast:UOp, opts:Renderer):

tinygrad/schedule/indexing.py

@@ -17,7 +17,7 @@ def realize_srcs(ctx:dict[UOp, None], rb:UOp) -> None:
     if s.base.op not in ALWAYS_CONTIGUOUS: ctx[s] = None
 
 def realize_assign(ctx:dict[UOp, None], a:UOp) -> None:
-  #if a.src[1].op not in ALWAYS_CONTIGUOUS: ctx[a.src[1]] = None
+  if a.src[1].op not in ALWAYS_CONTIGUOUS: ctx[a.src[1]] = None
   # if it's a kernel, we don't realize it
   if a.src[1].op is not Ops.KERNEL: ctx[a] = None
 
@@ -25,7 +25,7 @@ pm_generate_realize_map = PatternMatcher([
   # always realize SINK src
   (UPat(Ops.SINK, name="s"), lambda ctx,s: ctx.update((x.base, None) for x in s.src if x.base.op not in ALWAYS_CONTIGUOUS)),
   # always realize COPY/BUFFER_VIEW/CONTIGUOUS
-  (UPat({Ops.COPY, Ops.BUFFER_VIEW, Ops.CONTIGUOUS, Ops.ENDRANGE}, name="tr"), realize),
+  (UPat({Ops.COPY, Ops.BUFFER_VIEW, Ops.CONTIGUOUS}, name="tr"), realize),
   # realize srcs of COPY, MSELECT, MSTACK
   (UPat((Ops.COPY, Ops.MSELECT, Ops.MSTACK), name="rb"), realize_srcs),
   # realize ASSIGN and input to assign (might be optimized out)

tinygrad/schedule/rangeify.py

@@ -110,7 +110,7 @@ pm_mops = PatternMatcher([
 # 3.5 cleanups
 
 # Ops.NOOP happens when we have a COPY to the device the Tensor is already on. We treat it like COPY here for MSTACK.
-ALWAYS_RUN_OPS = {Ops.CONTIGUOUS, Ops.COPY, Ops.ASSIGN, Ops.NOOP, Ops.ENDRANGE}
+ALWAYS_RUN_OPS = {Ops.CONTIGUOUS, Ops.COPY, Ops.ASSIGN, Ops.NOOP}
 
 # you don't know in the first pass if axes are going to die, this happens if there's an EXPAND to the left
 def cleanup_dead_axes(b:UOp):
@@ -338,7 +338,6 @@ def handle_assign(ctx:LocalAddBufferContext, assign:UOp):
 
 def renumber_range(ctx:LocalAddBufferContext, r:UOp):
   if r.tag is not None: return None
-  if r.arg[-1] is AxisType.OUTER: return None
   ret = r.replace(arg=(ctx.range,)+r.arg[1:], tag=())
   ctx.range += 1
   return ret
@@ -413,7 +412,7 @@ class Kernel:
     return f"<Kernel {len(list(self.ast.toposort()))} {ast_rep} {self.metadata}>"
 
 def split_store(ctx:list[UOp], x:UOp):
-  if len([r for r in x.ranges if r.arg[-1] != AxisType.OUTER]): return None
+  if len(x.ranges): return None
   if x.src[0].ptrdtype.addrspace is AddrSpace.LOCAL: return None
 
   # local kernel rewrite
@@ -425,7 +424,7 @@ def split_store(ctx:list[UOp], x:UOp):
 
   # NOTE: the hack for COPY is here
   ret = ret.sink(arg=KernelInfo(opts_to_apply=lctx.opts) if lctx.opts is not None else None) \
-    if ret.src[1].op not in {Ops.COPY, Ops.BUFFER_VIEW, Ops.ENDRANGE} else ret.src[1]
+    if ret.src[1].op not in {Ops.COPY, Ops.BUFFER_VIEW} else ret.src[1]
   kernel_arg = Kernel(ret,tuple(dedup(flatten([x for x in metadatas if x is not None])))[::-1])
   kernel = UOp(Ops.KERNEL, src=tuple(lctx.map.values())+tuple(lctx.vars.keys()), arg=kernel_arg)
   if ret.op is Ops.SINK and not all_same([x.device for x in kernel.src if x.op is not Ops.BIND]):
@@ -482,11 +481,6 @@ def do_sub_recurse(s:UOp):
   return x.replace(src=tuple([UOp(Ops.SUBSTITUTE, dtype=y.dtype, src=(y,uop_keys,uop_values)) for y in x.src]))
 pm_substitute_recurse = PatternMatcher([(UPat(Ops.SUBSTITUTE, src=(UPat(), UPat(Ops.NOOP), UPat(Ops.NOOP)), name="s"), do_sub_recurse)])
 
-pm_localize_bufs = PatternMatcher([
-  (UPat(Ops.BUFFERIZE, name="x"), lambda x:
-   x.replace(arg=BufferizeOpts(device=None, addrspace=AddrSpace.LOCAL), tag=None) if len(x.ranges) > 0 else None),
-])
-
 @track_rewrites(lambda _,ret: f"Schedule {pluralize('Kernel', len([u for u in UOp.sink(*ret.values()).toposort() if u.op is Ops.KERNEL]))}", True)
 def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
   uop_list: list[UOp] = []
@@ -502,7 +496,7 @@ def get_rangeify_map(sink:UOp) -> dict[UOp, UOp]:
   tsink = graph_rewrite(tsink, pm_cleanups, bottom_up=True, name="remove costly buffers")
   # TODO: can you substitute and remove costly buffers at the same time?
   tsink = graph_rewrite(tsink, pm_substitute_recurse, bottom_up=True, name="run substitutes")
-  tsink = graph_rewrite(tsink, pm_localize_bufs+pm_limit_bufs, ctx=rctx, name="localize/limit buffers")
+  tsink = graph_rewrite(tsink, pm_limit_bufs, ctx=rctx, name="limit buffers")
 
   # rebuild the sink with all the BUFFERIZEs with tags, this is what's ending up in the tensor graph
   # MSTACK stacks multiple BUFFERIZEs in one tagged tensor

tinygrad/tensor.py

@@ -1218,7 +1218,7 @@ class Tensor(MathTrait):
           if not dtypes.is_int((ti:=Tensor(index)).dtype): raise IndexError(f"{index=} contains non-int element")
           index = Tensor([i+size if i<0 else i for i in fully_flatten(index)], self.device, requires_grad=False).reshape(ti.shape)
         case int() | UOp(): # sint
-          #if index >= size or index < -size: raise IndexError(f"{index=} is out of bounds with {size=}")
+          if index >= size or index < -size: raise IndexError(f"{index=} is out of bounds with {size=}")
           # TODO: is this right for (negative) symbolic?
           boundary = [index, index+1] if index >= 0 else [index+size, index+size+1]
         case slice():

tinygrad/uop/ops.py

@@ -15,7 +15,6 @@ if TYPE_CHECKING:
 
 class AxisType(Enum):
   def __repr__(self): return str(self)
-  OUTER = auto()
   GLOBAL = auto(); WARP = auto(); LOCAL = auto(); LOOP = auto(); GROUP_REDUCE = auto(); REDUCE = auto(); UPCAST = auto(); UNROLL = auto() # noqa: E702
   THREAD = auto()
 
@@ -241,7 +240,6 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
         if s in ret: del ret[s]
     else:
       for s in self.src: ret.update(s.ranges)
-    if self.op is Ops.ENDRANGE: del ret[self.src[0]]
     return ret
 
   @property
@@ -1100,8 +1098,6 @@ pm_lower_index_dtype = PatternMatcher([
   (UPat((Ops.STORE, Ops.LOAD), src=(UPat(), UPat(), UPat().cast(dtypes.index)), allow_any_len=True, name="s"),
     lambda s: s.replace(src=s.src[:2]+tuple(u.src[0] for u in s.src[2:]))),
   (UPat((Ops.SINK, Ops.NOOP), src=UPat().cast(dtypes.index), name="n"), lambda n: n.replace(src=tuple(s.src[0] for s in n.src))),
-  # hack for ENDRANGE
-  (UPat(Ops.ENDRANGE, src=(UPat(Ops.RANGE, name="r").cast(dtypes.index),), allow_any_len=True, name="x"), lambda x,r: x.replace(src=(r,)+x.src[1:])),
 ])
 def _index_to_concrete_int(u:UOp): return graph_rewrite(u.sink(), pm_lower_index_dtype).src[0]
 

tinygrad/uop/spec.py

@@ -109,9 +109,6 @@ tensor_uop_spec = buffer_spec+assign_spec+PatternMatcher([
   (UPat(Ops.ALLREDUCE, name="red", src=(UPat.var("x"), UPat(Ops.DEVICE))), lambda red,x: red.dtype == x.dtype and isinstance(red.arg, Ops)),
   (UPat(Ops.MULTI, name="multi"), lambda multi: all(x.dtype == multi.dtype for x in multi.src) and isinstance(multi.arg, int)),
 
-  # endrange/reduce for outerworld range work
-  (UPat(Ops.ENDRANGE, src=(UPat(Ops.RANGE),), allow_any_len=True), lambda: True),
-
   # REDUCE with an outerworld range
   (UPat(Ops.REDUCE, src=(UPat(),), allow_any_len=True, name="x"), lambda x: all(y.dtype == dtypes.index for y in x.src[1:])),
 ])

tinygrad/viz/js/index.js

@@ -18,9 +18,6 @@ const ANSI_COLORS_LIGHT = ["#d9d9d9","#ff9999","#99cc99","#ffff99","#9999ff","#f
 const parseColors = (name, defaultColor="#ffffff") => Array.from(name.matchAll(/(?:\u001b\[(\d+)m([\s\S]*?)\u001b\[0m)|([^\u001b]+)/g),
   ([_, code, colored_st, st]) => ({ st: colored_st ?? st, color: code != null ? (code>=90 ? ANSI_COLORS_LIGHT : ANSI_COLORS)[(parseInt(code)-30+60)%60] : defaultColor }));
 
-const colored = n => d3.create("span").call(s => s.selectAll("span").data(typeof n === "string" ? parseColors(n) : n).join("span")
-                       .style("color", d => d.color).text(d => d.st)).node();
-
 const rect = (s) => (typeof s === "string" ? document.querySelector(s) : s).getBoundingClientRect();
 
 let timeout = null;
@@ -177,7 +174,7 @@ function tabulate(rows) {
 var data, focusedDevice, focusedShape, canvasZoom, zoomLevel = d3.zoomIdentity;
 async function renderProfiler() {
   displayGraph("profiler");
-  d3.select(".metadata").node().replaceChildren(focusedShape?.html ?? "");
+  d3.select(".metadata").html("");
   // layout once!
   if (data != null) return updateProgress({ start:false });
   const profiler = d3.select(".profiler").html("");
@@ -239,7 +236,8 @@ async function renderProfiler() {
           const stepIdx = ctxs[ref.ctx+1].steps.findIndex((s, i) => i >= start && s.name == e.name);
           if (stepIdx !== -1) { ref.step = stepIdx; shapeRef = ref; }
         }
-        const arg = { tooltipText:colored(e.name).outerHTML+"\n"+formatTime(e.dur)+(e.info != null ? "\n"+e.info : ""), ...shapeRef };
+        const htmlLabel = label.map(({color, st}) => `<span style="color:${color}">${st}</span>`).join('');
+        const arg = { tooltipText:htmlLabel+"\n"+formatTime(e.dur)+(e.info != null ? "\n"+e.info : ""), ...shapeRef };
         // offset y by depth
         shapes.push({x:e.st, y:levelHeight*depth, width:e.dur, height:levelHeight, arg, label, fillColor });
       }
@@ -354,7 +352,7 @@ async function renderProfiler() {
           for (let i=x.length-1; i>=0; i--) p.lineTo(x[i], offsetY+e.y1[i]);
           p.closePath();
           ctx.fillStyle = e.fillColor; ctx.fill(p);
-          if (focusedShape && e.arg?.key === focusedShape.key) { paths.push(p); }
+          if (focusedShape && e.arg?.key === focusedShape) { paths.push(p); }
           continue;
         }
         // contiguous rect
@@ -450,7 +448,7 @@ async function renderProfiler() {
     e.preventDefault();
     const foundRect = findRectAtPosition(e.clientX, e.clientY);
     if (foundRect?.step != null) return setCtxWithHistory(foundRect.ctx, foundRect.step);
-    if (foundRect?.key != focusedShape?.key) { focusedShape = foundRect; render(zoomLevel); }
+    if (foundRect?.key != focusedShape) { focusedShape = foundRect?.key; render(zoomLevel); }
     return document.querySelector(".metadata").replaceChildren(foundRect?.html ?? "");
   });
 
@@ -594,7 +592,7 @@ async function main() {
       const ul = ctxList.appendChild(document.createElement("ul"));
       ul.id = `ctx-${i}`;
       const p = ul.appendChild(document.createElement("p"));
-      p.appendChild(colored(name));
+      p.innerHTML = parseColors(name).map(c => `<span style="color: ${c.color}">${c.st}</span>`).join("");
       p.onclick = () => {
         setState(i === state.currentCtx ? { expandSteps:!state.expandSteps } : { expandSteps:true, currentCtx:i, currentStep:0, currentRewrite:0 });
       }
@@ -708,7 +706,9 @@ async function main() {
         metadata.appendChild(codeBlock(upat[1], "python", { loc:upat[0], wrap:true }));
         const diffCode = metadata.appendChild(document.createElement("pre")).appendChild(document.createElement("code"));
         for (const line of diff) {
-          diffCode.appendChild(colored([{st:line, color:line.startsWith("+") ? "#3aa56d" : line.startsWith("-") ? "#d14b4b" : "#f0f0f5"}]));
+          const span = diffCode.appendChild(document.createElement("span"));
+          span.style.color = line.startsWith("+") ? "#3aa56d" : line.startsWith("-") ? "#d14b4b" : "#f0f0f5";
+          span.innerText = line;
           diffCode.appendChild(document.createElement("br"));
         }
         diffCode.className = "wrap";