dsp skip

test/unit/test_dtype_weak.py

@@ -11,7 +11,6 @@ class TestWeakPromotion(unittest.TestCase):
   def test_rand_requires_concrete(self):
     with self.assertRaises(ValueError): Tensor.rand(2, dtype=dtypes.weakfloat)
     with self.assertRaises(ValueError): Tensor.const(dtypes.weakfloat, 1.0).rand_like()
-    with self.assertRaises(ValueError): Tensor.const(dtypes.weakfloat, 1.0).randn_like()
 
   def test_sum_stays_weak(self):
     for weak, value in ((dtypes.weakint, 1), (dtypes.weakfloat, 1.0)):
@@ -23,17 +22,6 @@ class TestWeakPromotion(unittest.TestCase):
     for fn in (lambda: t.bitcast(dtypes.int32), lambda: Tensor.const(dtypes.int32, 2).bitcast(dtypes.weakint), t.element_size, t.nbytes):
       with self.assertRaises(RuntimeError): fn()
 
-  def test_materialize_at_default_dtype(self):
-    for weak, value, strong in ((dtypes.weakint, 3, dtypes.default_int), (dtypes.weakfloat, 0.5, dtypes.default_float)):
-      t = Tensor.const(weak, value)
-      self.assertEqual(t.dtype, weak)
-      self.assertEqual(t.data().itemsize, strong.itemsize)
-      self.assertEqual(t.numpy().dtype.itemsize, strong.itemsize)
-      realized = t.clone("CPU").realize()
-      self.assertEqual((realized.dtype, realized.uop.buffer.dtype), (strong, strong))
-    with patch.object(dtypes, "default_int", dtypes.int64):
-      self.assertEqual(Tensor.const(dtypes.weakint, 3).numpy().dtype.itemsize, dtypes.int64.itemsize)
-
   def test_uop_scalar_const_unchanged(self):
     for dtype, value in ((dtypes.index, 1), (dtypes.int32, 1), (dtypes.float32, 0.5)):
       out = UOp.variable("x", 0.0 if dtype == dtypes.float32 else 0, 10.0 if dtype == dtypes.float32 else 10, dtype) + value

test/unit/test_tensor_data.py

@@ -83,10 +83,5 @@ class TestTensorData(unittest.TestCase):
     assert dat.shape == (2,2)
     # NOTE: python can't deref float16
 
-  def test_tolist_empty_shapes(self):
-    for shape, expected in (((0,), []), ((2, 0), [[], []]), ((0, 2), []),
-                            ((2, 0, 3), [[], []]), ((2, 3, 0), [[[], [], []], [[], [], []]])):
-      self.assertEqual(Tensor.ones(*shape).tolist(), expected)
-
 if __name__ == '__main__':
   unittest.main()

tinygrad/codegen/__init__.py

@@ -79,10 +79,9 @@ def unroll_axis(ctx:dict[int, int], u:UOp, arg):
   return out.permute(argsort(permute_head+permute_tail))
 
 def expand_wmma(ctx:dict[int, int], u:UOp):
-  if u.arg[4] is None: return None
+  if u.tag != 1: return None
   in0, in1, out0 = u.arg[4]
-  wmma = u.replace(src=(contract_axis(ctx, u.src[0], in0), contract_axis(ctx, u.src[1], in1), u.src[2]),
-                   arg=(*u.arg[:4], None))
+  wmma = u.replace(src=(contract_axis(ctx, u.src[0], in0), contract_axis(ctx, u.src[1], in1), u.src[2]), tag=None)
   return unroll_axis(ctx, wmma, out0)
 
 expander2 = PatternMatcher([

tinygrad/codegen/opt/postrange.py

@@ -302,7 +302,7 @@ class Scheduler:
             # do the reduce_axes always disappear? i think they don't
             # they need to be moved into the WMMA srcs
             tc_uop = UOp.wmma(srcs[0], srcs[1], UOp.const(tc.dtype_out, (0.0,)*tc.elements_per_thread[2]),
-                              tc.dims, self.ren.target.device, tc.threads, tc_upcast_axes=tc_upcast_axes)
+                              tc.dims, self.ren.target.device, tc.threads, tag=1, tc_upcast_axes=tc_upcast_axes)
 
             # preserve extra reduces
             reduce_ranges = [x for x in UOp.sink(*reduceop.src[1:]).toposort() if x.op is Ops.RANGE and x.arg[0] not in tc_reduce_axes]

tinygrad/dtype.py

@@ -163,8 +163,6 @@ if (env_default_float := getenv("DEFAULT_FLOAT", "")):
 
 DTypeLike = str|DType
 def to_dtype(dtype:DTypeLike) -> DType: return dtype if isinstance(dtype, DType) else getattr(dtypes, dtype.lower())
-def strong_dtype(dtype:DType) -> DType:
-  return dtypes.default_int if dtype == dtypes.weakint else dtypes.default_float if dtype == dtypes.weakfloat else dtype
 
 # https://jax.readthedocs.io/en/latest/jep/9407-type-promotion.html
 # we don't support complex type

tinygrad/llm/cli.py

@@ -6,26 +6,6 @@ from tinygrad.helpers import partition, DEBUG, Timing, GlobalCounters, stderr_lo
 from tinygrad.viz.serve import TCPServerWithReuse, HTTPRequestHandler
 from tinygrad.llm.model import Transformer
 
-def holdback(s:str, tag:str) -> int:
-  # length of the suffix of s that is a prefix of tag (the tag may be split across streamed pieces)
-  return max((i for i in range(1, min(len(s), len(tag))+1) if tag.startswith(s[-i:])), default=0)
-
-def parse_tool_call(s:str) -> tuple[str, typing.Any]|None:
-  s = s.strip()
-  if s.startswith("{"):  # hermes JSON format: {"name": ..., "arguments": {...}}
-    try:
-      call = json.loads(s)
-      return call["name"], call.get("arguments", call.get("parameters", {}))
-    except (json.JSONDecodeError, KeyError): return None
-  # XML format: <function=name>\n<parameter=key>\nvalue\n</parameter>...</function>
-  if (fm := re.match(r"<function=([^>]+)>\s*(.*?)\s*(?:</function>)?$", s, re.DOTALL)):
-    args = {}
-    for pm in re.finditer(r"<parameter=([^>]+)>\s*(.*?)\s*</parameter>", fm.group(2), re.DOTALL):
-      try: args[pm.group(1)] = json.loads(pm.group(2))
-      except json.JSONDecodeError: args[pm.group(1)] = pm.group(2)
-    return fm.group(1), args
-  return None
-
 class SimpleTokenizer:
   def __init__(self, normal_tokens:dict[str, int], special_tokens:dict[str, int], preset:str="llama3",
                bos_id:int|None=None, eos_id:int=0, eot_id:int|None=None):
@@ -133,75 +113,26 @@ class Handler(HTTPRequestHandler):
   def do_GET(self):
     if self.path == "/v1/models": self.send_data(json.dumps({"object":"list","data":[{"id":self.server.model_name,"object":"model"}]}).encode())
     else: self.send_data((pathlib.Path(__file__).parent / "chat.html").read_bytes(), content_type="text/html")
-  def run_model(self, ids:list[int], model_name:str, include_usage=False, max_tokens:int|None=None, temperature:float=0.0,
-                parse_tool_calls=False, prefill_think=False):
+  def run_model(self, ids:list[int], model_name:str, include_usage=False, max_tokens:int|None=None, temperature:float=0.0):
     model, tok = self.server.model, self.server.tok
     cache_start_pos = model.get_start_pos(ids)
     stderr_log(f"{self.path}  {colored('--', 'BLACK')}  "
                f"in:{colored(f'{cache_start_pos:5d}', 'green')} +{len(ids)-cache_start_pos:5d}  {colored('--', 'BLACK')}  ")
     tmpl = {"id":f"chatcmpl-{uuid.uuid4().hex[:24]}", "object":"chat.completion.chunk", "created":int(time.time()), "model":model_name}
-    def chunk(d:dict): return {"choices": [{"index":0, "delta":d, "finish_reason":None}], **tmpl}
-    yield chunk({"role":"assistant", "content":""})
+    yield {"choices": [{"index":0, "delta":{"role":"assistant","content":""}, "finish_reason":None}], **tmpl}
     out: list[int] = []
     finish_reason = "stop"
     st = time.perf_counter()
     dec = tok.stream_decoder()
-    mode, buf, tool_text = ("reasoning" if prefill_think else "undecided"), "", ""
-    def route(piece:str, final:bool=False):
-      nonlocal mode, buf, tool_text
-      if mode == "undecided":  # decide whether the output starts with a think block
-        buf += piece
-        if not final and len(buf) < len("<think>") and "<think>".startswith(buf): return
-        mode, piece, buf = ("reasoning", buf[len("<think>"):], "") if buf.startswith("<think>") else ("content", buf, "")
-      if mode == "reasoning":
-        buf += piece
-        if "</think>" in buf:
-          before, piece = buf.split("</think>", 1)
-          if before: yield chunk({"reasoning_content":before})
-          buf, mode, piece = "", "content", piece.lstrip("\n")
-        else:
-          hold = 0 if final else holdback(buf, "</think>")
-          if (flush := buf[:len(buf)-hold]): yield chunk({"reasoning_content":flush})
-          buf = buf[len(buf)-hold:]
-          return
-      if not parse_tool_calls:
-        if piece: yield chunk({"content":piece})
-      else:
-        tool_text += piece
-        if tool_text.startswith("<tool_call>"): return
-        if "<tool_call>" in tool_text:
-          before, tool_text = tool_text.split("<tool_call>", 1)
-          if before: yield chunk({"content":before})
-          tool_text = "<tool_call>" + tool_text
-        else:
-          # hold back any suffix that could be the start of a "<tool_call>" tag split across tokens
-          hold = 0 if final else holdback(tool_text, "<tool_call>")
-          if (flush := tool_text[:len(tool_text)-hold]): yield chunk({"content":flush})
-          tool_text = tool_text[len(tool_text)-hold:]
     for next_id in model.generate(ids, temperature=temperature):
       if len(out) == 0: stderr_log(f"prefill:{(len(ids)-cache_start_pos)/((pt:=time.perf_counter())-st):4.0f} tok/s  {colored('--', 'BLACK')}  ")
       if tok.is_end(next_id): break
       out.append(next_id)
-      yield from route(dec(next_id))
+      yield {"choices": [{"index":0, "delta":{"content":dec(next_id)}, "finish_reason":None}], **tmpl}
       if max_tokens is not None and len(out) >= max_tokens:
         finish_reason = "length"
         break
-    yield from route(dec(), final=True)
-    if parse_tool_calls:
-      tool_calls = []
-      calls = [(m.group(1), m.group(0)) for m in re.finditer(r"<tool_call>\s*(.*?)\s*</tool_call>", tool_text, re.DOTALL)]
-      if not calls and tool_text.startswith("<tool_call>"): calls = [(tool_text[len("<tool_call>"):], tool_text)]  # unclosed tag
-      for i, (inner, raw) in enumerate(calls):
-        if (parsed := parse_tool_call(inner)) is None:
-          stderr_log(f"failed to parse tool call: {inner[:200]}")
-          yield chunk({"content":raw})  # don't silently drop output the client can't use
-        else:
-          name, args = parsed
-          tool_calls.append({"index":i, "id":f"call_{uuid.uuid4().hex[:24]}", "type":"function",
-                             "function":{"name":name, "arguments":args if isinstance(args, str) else json.dumps(args)}})
-      if tool_calls:
-        yield chunk({"tool_calls":tool_calls})
-        if finish_reason == "stop": finish_reason = "tool_calls"
+    if (tail := dec()): yield {"choices": [{"index":0, "delta":{"content":tail}, "finish_reason":None}], **tmpl}
     yield {"choices": [{"index":0, "delta":{},"finish_reason":finish_reason}], **tmpl}
     if include_usage:
       yield {"choices": [], "usage": {"prompt_tokens": len(ids), "completion_tokens": len(out), "total_tokens": len(ids) + len(out)}, **tmpl}
@@ -215,65 +146,39 @@ class Handler(HTTPRequestHandler):
     body: dict[str, typing.Any] = json.loads(raw_body.decode("utf-8"))
     if DEBUG >= 1: print(json.dumps(body, indent=2))
     if self.path == "/v1/chat/completions":
-      messages, tools = body["messages"], body.get("tools")
-      prefill_think = False
-      if self.server.template is not None:
-        # the chat template expects tool_call arguments as dicts, OpenAI clients send them as JSON strings
-        norm = []
-        for m in messages:
-          if m.get("tool_calls"):
-            m = dict(m)
-            m["tool_calls"] = [{**tc, "function":{**tc["function"], "arguments":json.loads(a) if isinstance((a:=tc["function"]["arguments"]), str)
-                                else a}} if "function" in tc else tc for tc in m["tool_calls"]]
-          norm.append(m)
-        rendered = self.server.template.render(messages=norm, tools=tools, add_generation_prompt=norm[-1]["role"] != "assistant")
-        prefill_think = rendered.rstrip().endswith("<think>")
-        ids: list[int] = tok.encode(rendered)
-        if (prefix := tok.prefix()) and ids[:len(prefix)] != prefix: ids = prefix + ids
-        if norm[-1]["role"] == "assistant":  # last assistant message is treated as prefill, drop its end-of-turn tokens
-          end = tok.end_turn()
-          if len(ids) >= len(end) and ids[-len(end):] == end: ids = ids[:-len(end)]
-      else:
-        if tools: stderr_log("warning: ignoring tools, install jinja2 to enable tool calling via the model's chat template")
-        ids = tok.prefix()
-        for i, msg in enumerate(messages):
-          ids += tok.role(msg["role"])
-          content = msg["content"]
-          if isinstance(content, str): ids += tok.encode(content)
-          elif isinstance(content, list):
-            for c in content:
-              if c["type"] == "text": ids += tok.encode(c["text"])
-              else: raise RuntimeError(f"unhandled type: {c['type']}")
-          else: raise RuntimeError(f"unknown content type: {type(content)}")
-          if msg["role"] == "assistant" and i == len(messages) - 1: break
-          ids += tok.end_turn()
-        else: ids += tok.role("assistant")
+      # extract tokens, last assistant message is treated as prefill
+      ids: list[int] = tok.prefix()
+      for i, msg in enumerate(body["messages"]):
+        ids += tok.role(msg["role"])
+        content = msg["content"]
+        if isinstance(content, str): ids += tok.encode(content)
+        elif isinstance(content, list):
+          for c in content:
+            if c["type"] == "text": ids += tok.encode(c["text"])
+            else: raise RuntimeError(f"unhandled type: {c['type']}")
+        else: raise RuntimeError(f"unknown content type: {type(content)}")
+        if msg["role"] == "assistant" and i == len(body["messages"]) - 1: break
+        ids += tok.end_turn()
+      else: ids += tok.role("assistant")
 
       # reply
       max_tokens = body.get("max_completion_tokens") or body.get("max_tokens")
       chunks = self.run_model(ids, body["model"], not body.get("stream") or body.get("stream_options",{}).get("include_usage", False),
-                              max_tokens=max_tokens, temperature=float(body.get("temperature", 0.0)), parse_tool_calls=bool(tools),
-                              prefill_think=prefill_think)
+                              max_tokens=max_tokens, temperature=float(body.get("temperature", 0.0)))
       if body.get("stream"): self.stream_json(chunks)
       else:
-        out, reasoning, tool_calls, finish_reason = [], [], [], "stop"
+        out, finish_reason = [], "stop"
         for c in chunks:
-          if c["choices"] and (delta := c["choices"][0].get("delta", {})):
-            if delta.get("content"): out.append(delta["content"])
-            if delta.get("reasoning_content"): reasoning.append(delta["reasoning_content"])
-            if delta.get("tool_calls"): tool_calls.extend(delta["tool_calls"])
+          if c["choices"] and c["choices"][0].get("delta", {}).get("content"): out.append(c["choices"][0]["delta"]["content"])
           if c["choices"] and c["choices"][0].get("finish_reason"): finish_reason = c["choices"][0]["finish_reason"]
-        message: dict[str, typing.Any] = {"role":"assistant", "content":"".join(out) or None}
-        if reasoning: message["reasoning_content"] = "".join(reasoning)
-        if tool_calls: message["tool_calls"] = [{k:v for k, v in tc.items() if k != "index"} for tc in tool_calls]
         self.send_data(json.dumps({**c, "object":"chat.completion",
-          "choices":[{"index":0, "message":message, "finish_reason":finish_reason}]}).encode())
+          "choices":[{"index":0, "message":{"role":"assistant","content":"".join(out)}, "finish_reason":finish_reason}]}).encode())
     else:
       raise RuntimeError(f"unhandled path {self.path}")
 
 class LLMServer(TCPServerWithReuse):
-  def __init__(self, server_address:tuple, model:Transformer, model_name:str, tok:SimpleTokenizer, template:typing.Any=None):
-    self.model, self.model_name, self.tok, self.template = model, model_name, tok, template
+  def __init__(self, server_address:tuple, model:Transformer, model_name:str, tok:SimpleTokenizer):
+    self.model, self.model_name, self.tok = model, model_name, tok
     super().__init__(server_address, Handler)
 
 def main():
@@ -289,26 +194,11 @@ def main():
   model, kv = Transformer.from_gguf(fetch(models.get(args.model, args.model)), args.max_context)
   model_name = kv.get('general.name') or kv.get('general.basename') or args.model
   file_sizes = [y.nbytes() for y in UOp.sink(*[x.uop for x in nn.state.get_parameters(model)]).toposort() if y.op is Ops.BUFFER]
-  print(f"using model \"{model_name}\" with {sum(file_sizes):,} bytes and {sum(x.numel() for x in nn.state.get_parameters(model)):,} params, "
-        f"max context {args.max_context} on {nn.state.get_parameters(model)[0].device}")
+  print(f"using model \"{model_name}\" with {sum(file_sizes):,} bytes and {sum(x.numel() for x in nn.state.get_parameters(model)):,} params")
 
   # get tokenizer
   tok = SimpleTokenizer.from_gguf_kv(kv)
 
-  # compile the chat template if jinja2 is available (enables tool calling and model-specific formatting)
-  template = None
-  if (ct := kv.get('tokenizer.chat_template')) is not None:
-    try:
-      import jinja2
-      env = jinja2.Environment()
-      env.filters['tojson'] = lambda obj, **kwargs: json.dumps(obj)  # jinja2's tojson escapes <>& for HTML safety
-      env.globals['raise_exception'] = lambda msg: (_ for _ in ()).throw(RuntimeError(msg))
-      env.globals['strftime_now'] = lambda fmt: time.strftime(fmt)
-      for name, key in {'bos_token':'bos', 'eos_token':'eos', 'unk_token':'unknown', 'pad_token':'padding', 'sep_token':'separator'}.items():
-        if (tid := kv.get(f'tokenizer.ggml.{key}_token_id')) is not None: env.globals[name] = tok.decode([tid])
-      template = env.from_string(ct)
-    except ImportError: stderr_log("warning: jinja2 is not installed, the model's chat template is disabled")
-
   # warmup the JIT
   if args.warmup or args.serve:
     # run 2 tokens through the model twice to capture the JIT before serving
@@ -316,7 +206,7 @@ def main():
       for _ in range(2): list(zip(range(2), model.generate([0])))
 
   # start server
-  if args.serve: LLMServer(('', args.serve), model, model_name, tok, template).serve_forever()
+  if args.serve: LLMServer(('', args.serve), model, model_name, tok).serve_forever()
 
   # do benchmark
   if args.benchmark is not None:

tinygrad/mixin/rand.py

@@ -97,10 +97,9 @@ class RandMixin(OpMixin):
     print(Tensor.randn_like(t).numpy())
     ```
     """
-    if (dt:=to_dtype(dtype or self.dtype)) in dtypes.weaks and dtype is None: raise ValueError(f"randn_like requires an explicit dtype for {dt}")
     src = self.stack(self).rand_like(**{**kwargs, "dtype": dtypes.float32})
     # https://en.wikipedia.org/wiki/Box%E2%80%93Muller_transform
-    return src[0].mul(2*math.pi).cos().mul((1 - src[1]).log().mul(-2).sqrt()).cast(dt)
+    return src[0].mul(2*math.pi).cos().mul((1 - src[1]).log().mul(-2).sqrt()).cast(to_dtype(dtype or self.dtype))
 
   @classmethod
   def randn(cls, *shape, dtype:DTypeLike|None=None, **kwargs) -> Self:

tinygrad/renderer/cstyle.py

@@ -101,11 +101,10 @@ def uops_to_dtypes(uops:list[UOp]) -> list[tuple[DType, int]]:
 def _wmma_name(u:UOp) -> str:
   return f"WMMA_{'_'.join(map(str, u.arg[0]))}_{u.arg[1].name}_{u.dtype.scalar().name}"
 
-# (name, dims, dtype_in, dtype_out, device, threads, upcast_sizes)
+# (name, dims, dtype_in, dtype_out, device, threads, upcast_axes)
 def wmma_args(uops:list[UOp]):
-  return dedup((_wmma_name(uop), uop.arg[0], uop.arg[1], uop.dtype.scalar(), *(uop.arg[2:4]),
-               tuple(uop.src[i].shape[-1] for i in range(3)))
-              for uop in uops if uop.op is Ops.WMMA)
+  return dedup((_wmma_name(uop), uop.arg[0], uop.arg[1], uop.dtype.scalar(), *(uop.arg[2:5]))
+                for uop in uops if uop.op is Ops.WMMA)
 
 class CStyleLanguage(Renderer):
   kernel_typedef: str = "void"
@@ -443,7 +442,8 @@ class CUDARenderer(CStyleLanguage):
       or (count in (2,4,8,16) and dt in dtypes.fp8s)]
     dt_map_in = { dtypes.float: "tf32", dtypes.half: "f16", dtypes.bfloat16: "bf16", dtypes.fp8e4m3: "e4m3", dtypes.fp8e5m2: "e5m2" }
     dt_map_out = { dtypes.float: "f32", dtypes.half: "f16" }
-    for name, (N, M, K), dtype_in, dtype_out, _, _, upcast_sizes in wmma_args(uops):
+    for name, (N, M, K), dtype_in, dtype_out, _, _, upcast_axes in wmma_args(uops):
+      upcast_sizes = [prod(size for _, size in upcast) for upcast in upcast_axes]
       wmma_dtypes = [self._render_dtype(dtype, size, AddrSpace.REG) for dtype, size in zip([dtype_in, dtype_in, dtype_out], upcast_sizes)]
       n_operands = [size*dtype.itemsize//4 for dtype, size in zip([dtype_in, dtype_in, dtype_out], upcast_sizes)] # 4 => CUDA reg size in bytes
       operands = [f"%{i}" for i in range(sum(n_operands))]

tinygrad/renderer/llvmir.py

@@ -264,7 +264,7 @@ exit: %packed = phi i32 [%packed_bf8, %do_bf8], [%packed_fp8, %do_fp8]\n  %trunc
         (UPat(Ops.WMMA, name="x", dtype=dtypes.half), lambda x: UOp(Ops.STACK, src=tuple(x.replace(
           src=(x.src[0], x.src[1], UOp(Ops.STACK, src=tuple(x.src[2].index(j//2) if j%2 == 0 else UOp.const(x.src[2].dtype, 0.0)
             for j in range(x.max_numel()*2)))),
-          arg=(*x.arg[:4], None)).index(i*2)
+          arg=(*x.arg[:4], (*x.arg[4][:2], ((0, x.max_numel()*2),)))).index(i*2)
           for i in range(x.max_numel()))) if x.max_numel() == 8 else None),
         (UPat(Ops.WMMA, name="x"), lambda x: x.replace(
           src=(x.src[0].bitcast(dtypes.uint16), x.src[1].bitcast(dtypes.uint16), x.src[2]))

tinygrad/runtime/support/system.py

@@ -165,10 +165,6 @@ class PCIDevice:
       FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/driver/unbind", os.O_WRONLY).write(self.pcibus)
     if FileIOInterface.exists(f"/sys/bus/pci/devices/{self.pcibus}/driver"): raise RuntimeError(f"Driver is bound to {pcibus}")
 
-    # remove sibling functions of the gpu, if any
-    for fn in range(1, 8):
-      if FileIOInterface.exists(sib:=f"/sys/bus/pci/devices/{self.pcibus[:-1]}{fn}"): FileIOInterface(f"{sib}/remove", os.O_WRONLY).write("1")
-
     if getenv("VFIO", 0) and (vfio_fd:=System.vfio) is not None:
       FileIOInterface(f"/sys/bus/pci/devices/{self.pcibus}/driver_override", os.O_WRONLY).write("vfio-pci")
       FileIOInterface("/sys/bus/pci/drivers_probe", os.O_WRONLY).write(self.pcibus)

tinygrad/tensor.py

@@ -3,7 +3,7 @@ from __future__ import annotations
 import time, functools, sys, inspect, pathlib, hashlib, weakref
 from typing import Any, Callable, cast, get_args, ParamSpec, TypeVar, Generic, TYPE_CHECKING
 if TYPE_CHECKING: import numpy
-from tinygrad.dtype import DType, DTypeLike, dtypes, ConstType, to_dtype, strong_dtype, _from_np_dtype, _to_np_dtype, PyConst
+from tinygrad.dtype import DType, DTypeLike, dtypes, ConstType, to_dtype, _from_np_dtype, _to_np_dtype, PyConst
 from tinygrad.helpers import all_int, getenv, fully_flatten, fetch, Metadata, TRACEMETA, is_numpy_ndarray, TracingKey
 from tinygrad.helpers import cpu_profile, suppress_finalizing, disable_gc
 from tinygrad.uop.ops import UOp, Ops, sint, all_metadata, _index_to_concrete_int, Variable, _broadcast_shape
@@ -238,7 +238,7 @@ class Tensor(RandMixin):
     if capturing and not getenv("UNSAFE_ALLOW_JIT_BUFFER"):
       from tinygrad.engine.jit import JitError
       raise JitError("cannot access tensor data during JIT capture, the value will be baked in")
-    x = self.cast(strong_dtype(self.dtype)).contiguous()
+    x = self.cast(self.dtype).contiguous()
     if self.uop.device is None or isinstance(self.device, tuple): x = x.clone("CPU")
     return cast(Buffer, x.realize().uop.buffer).ensure_allocated()
 
@@ -255,11 +255,10 @@ class Tensor(RandMixin):
     """
     if 0 in self.shape: return memoryview(bytearray(0)).cast(self.dtype.fmt)  # type: ignore[arg-type,return-value]
     assert all_int(self.shape), f"no data if shape is symbolic, {self.shape=}"
-    buf = self._buffer()
-    fmt = buf.dtype.fmt
-    assert fmt is not None, f"no fmt dtype for {buf.dtype}"
+    fmt = self.dtype.fmt
+    assert fmt is not None, f"no fmt dtype for {self.dtype}"
     assert fmt != "e" or sys.version_info >= (3, 12)
-    return buf.as_memoryview().cast(fmt, self.shape)  # type: ignore[arg-type,return-value]
+    return self._data().cast(fmt, self.shape)  # type: ignore[arg-type,return-value]
 
   # NOTE: list[Any] because return type is recursive (list[list[...]] for higher dimensions)
   def tolist(self) -> PyConst|list[Any]:
@@ -278,10 +277,6 @@ class Tensor(RandMixin):
     """
     # TODO: remove half once minimum python supports it
     if self.dtype in (dtypes.half, dtypes.bfloat16, *dtypes.fp8s): return self.cast(dtypes.float32).tolist()
-    if 0 in self.shape:
-      assert all_int(self.shape), f"no data if shape is symbolic, {self.shape=}"
-      def _tolist(shape:tuple[int, ...]): return [_tolist(shape[1:]) for _ in range(shape[0])]
-      return _tolist(self.shape)
     return self.data().tolist()
 
   def numpy(self) -> 'numpy.ndarray':

tinygrad/uop/ops.py

@@ -4,7 +4,7 @@ import sys, time, functools, itertools, math, operator, hashlib, os, types, pick
 from dataclasses import dataclass, replace
 from enum import Enum, auto
 from tinygrad.uop import Ops, GroupOp
-from tinygrad.dtype import ConstType, dtypes, DType, DTypeLike, truncate, least_upper_dtype, least_upper_float, strong_dtype, Invalid, AddrSpace
+from tinygrad.dtype import ConstType, dtypes, DType, DTypeLike, truncate, least_upper_dtype, least_upper_float, Invalid, AddrSpace
 from tinygrad.dtype import ConstFloat, PyConst, InvalidType, storage_fmt_for_dtype, to_storage_scalar, from_storage_scalar
 from tinygrad.device import Buffer, MultiBuffer, canonicalize_device
 from tinygrad.helpers import ContextVar, all_int, prod, getenv, all_same, Context, partition, temp, unwrap, T, argfix, Metadata, flatten, TRACEMETA
@@ -367,8 +367,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
 
       # wmma output shape = accumulator shape (src[2])
       case Ops.WMMA:
+        in0, in1, out0 = self.arg[4]
         wmma_b = _broadcast_shape(self.src[0].shape[:-1], self.src[1].shape[:-1], self.src[2].shape[:-1])
-        return wmma_b + (self.src[2].shape[-1],)
+        return wmma_b + (prod([x for _,x in out0]),)
 
       # passthrough ops
       case Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.AFTER | Ops.LOAD | \
@@ -600,9 +601,11 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
   @staticmethod
   def special(end:sint, name:str, dtype=dtypes.index): return UOp(Ops.SPECIAL, src=(sint_to_uop(end, dtype),), arg=name)
   @staticmethod
-  def wmma(a:UOp, b:UOp, acc:UOp, dims:tuple[int, int, int], device:str, threads:int, tc_upcast_axes=None):
+  def wmma(a:UOp, b:UOp, acc:UOp, dims:tuple[int, int, int], device:str, threads:int, tag=None, tc_upcast_axes=None):
+    if tc_upcast_axes is None:
+      tc_upcast_axes = tuple(((i, s.shape[-1] if s.shape else 1),) for i,s in enumerate((a, b, acc)))
     # dtype_in is stored in the arg (not derived from src[0].dtype) because bitcast rewrites change src dtypes
-    return UOp(Ops.WMMA, src=(a, b, acc), arg=(dims, a.dtype, device, threads, tc_upcast_axes))
+    return UOp(Ops.WMMA, src=(a, b, acc), arg=(dims, a.dtype, device, threads, tc_upcast_axes), tag=tag)
   def _rop(self, op:Ops, axis:tuple[int, ...]):
     # NOTE: we don't allow reduce on 1s axis
     axis = tuple(sorted(axis))
@@ -785,9 +788,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
     return ret if ret.device == device else ret.copy_to_device(device)
   def clone(self, device=None) -> UOp:
     device = device or self.device
-    ret = self.empty_like(dtype=strong_dtype(self.dtype), device=device)
+    ret = self.empty_like(device=device)
     src = self if self.device is None or self.device == device else self.copy_to_device(device)
-    return ret.after(ret.store(src.cast(ret.dtype)))
+    return ret.after(ret.store(src))
   @recursive_property
   def device(self) -> str|tuple[str, ...]|None:
     if self.op is Ops.PARAM: return self.arg.device

tinygrad/uop/symbolic.py

@@ -347,7 +347,7 @@ def uop_given_valid(valid:UOp, uop:UOp, try_simplex=True) -> UOp:
 
 def _valid_priority(v: UOp, valids:list[UOp]) -> int:
   # we want valid that's in other valids' parents to be first, so it's more likely the other valids get simplified
-  return 0 if (res:=parse_valid(v)) is None else sum(-1 for other in valids if res[0] in other.backward_slice_with_self)
+  return sum(-1 if (res:=parse_valid(v)) is not None and res[0] in other.toposort() else 0 for other in valids)
 
 def simplify_valid(valid:UOp) -> UOp|None:
   if valid.op_in_backward_slice_with_self(Ops.INDEX): return None  # this should only be for indexing, skip if there's a INDEX