pad invalid work (glm)

* set tc_upcast_axes to None when done with it

* no tag needed

CONTINUE.md

@@ -0,0 +1,122 @@
+# CONTINUE.md: PAD with Invalid instead of 0
+
+## Goal
+Make the low-level `Ops.PAD` pad with `Invalid` instead of `0`, while keeping
+the external `Tensor.pad` behavior unchanged.
+
+## Changes made (all 3 files are modified, see `git diff`)
+
+### 1. `tinygrad/schedule/indexing.py:92` — core change
+`convert_pad_to_where_to_keep_behavior_local` now uses `UOp.const(x.dtype, Invalid)`
+instead of `UOp.const(x.dtype, 0)` as the else value. This is what makes `Ops.PAD`
+pad with Invalid.
+
+### 2. `tinygrad/uop/symbolic.py:87-99` — Invalid propagation rules
+Added two new rules to `pm_data_invalid` so that `where(invalid_gate, a, const_b)`
+uses `b` (the const) in don't-care positions instead of poisoning to Invalid.
+This is needed so that `_pad_constant`'s mask `where(pad(ones_bool), base, value)`
+works — the mask is a `where(valid, True, Invalid)` gate, and the else `value`
+is a const.
+
+The rules are restricted to only match when the gate's valid value is a **const**
+(`UPat.cvar("x")`), to distinguish pad masks (where valid=True, a const) from
+gather masks (where valid=loaded_data, not a const). Without this restriction,
+`test_tensor_index` breaks because gather masks also create `where(cond, x, Invalid)`
+but need to keep poisoning.
+
+### 3. `tinygrad/mixin/op.py:280-289` — `_pad_constant` fix
+Swapped the `value == 0` early return for `value is Invalid` early return.
+When `value is Invalid`, just return `base` (which already has Invalid from
+`Ops.PAD`). For all other values (including 0), use the mask approach:
+`where(pad(ones_bool), base, const_value)`.
+
+## Current state
+- `test/unit/test_invalid_tensor.py` — **all 22 pass**
+- `test/unit/test_function.py` — **5 failures**, all multi-shard tests
+
+## The remaining bug: `cat` + multi-shard
+
+`cat` (op.py:716) uses `pad` + `usum` (element-wise ADD) to combine tensors:
+```python
+padded = [t.pad(...) for i,t in enumerate(tensors)]
+return padded[0].usum(*padded[1:])
+```
+
+When two shards are cat'd, each is padded and then summed. The valid masks
+are **complementary** (shard 0 valid in positions 0-1, shard 1 valid in 2-3).
+
+`_pad_constant` creates `where(mask_pad, data_pad, 0)` where:
+- `mask_pad = where(valid, True, Invalid)` — gate's valid value is const `True`
+- `data_pad = where(valid, data, Invalid)` — gate's valid value is loaded `data` (NOT const)
+
+The new const-specific rule handles the mask pad correctly. But for the data pad,
+the gate's valid value (`data`) is not a const, so the **non-const** lift-out rule
+fires: `where(valid, where(valid, data, Invalid), 0)` → `where(valid, where(valid, data, 0), Invalid)`.
+
+The `Invalid` else poisons the ADD. The binary Invalid rule lifts both gates out:
+`where(c6, data0, Invalid) + where(c8, data1, Invalid)` → `where(c6&c8, data0+data1, Invalid)`.
+
+Since `c6` and `c8` are complementary, `c6&c8` is always False → result is all Invalid → 0.
+
+### Master comparison
+On master, `convert_pad_to_where` uses `0` (not Invalid), so the ADD is just
+`where(c6, data0, 0) + where(c8, data1, 0)` with no Invalid, no lifting, works fine.
+
+### Debug output (with changes)
+```
+c16 = c6.where(c11.index(c13), 0)        # where(c6, load0, 0) — correct
+c22 = c6.where(0, c17.index(c20))         # where(c6, 0, load1) — correct
+c25 = (c6&c8).where((c16+c22), Invalid)   # WRONG: c6&c8 always False → all Invalid
+```
+
+### Master debug output
+```
+c13 = c6.where(c8.index(c10), 0)          # where(c6, load0, 0)
+c21 = c6.where(0, c14.index(c19))          # where(c6, 0, load1)
+c22 = c13+c21                              # plain ADD, no wrapper — correct
+```
+
+## Suggested fix approaches
+
+### Option A: General WHERE simplification rule
+Add a rule: `where(a, where(a, x, _), c)` → `where(a, x, c)`.
+When the outer and inner conditions are the same UOp, the inner else is
+unreachable. This would simplify `where(valid, where(valid, data, Invalid), 0)`
+→ `where(valid, data, 0)` before the lift-out rule can fire.
+Check if this rule already exists in `symbolic.py` — it may need to be added
+before the lift-out rules.
+
+### Option B: Don't use Ops.PAD for data in `_pad_constant`
+When `value is not Invalid`, avoid creating `Ops.PAD` on the data. Use `cat`
+or `expand` to create the padded tensor directly, bypassing the Invalid
+propagation entirely.
+
+### Option C: Make the lift-out rule use the outer else value
+Change the non-const lift-out rule: when `where(a, where(cond, x, Invalid), c)`
+and `c` is a const, use `c` as the else instead of `Invalid`. This is what the
+const-specific rule does, but it needs to also handle non-const gate valid values.
+
+## Test commands
+```bash
+# invalid tensor tests (currently pass)
+python -m pytest test/unit/test_invalid_tensor.py -x -q -n12
+
+# function tests (5 multi-shard failures)
+python -m pytest test/unit/test_function.py -x -q -n12
+
+# the specific failing test
+python -m pytest test/unit/test_function.py::TestFunctionMulti::test_simple_multi_sharded -x -q
+
+# debug the failing case
+DEBUG=6 python -c "
+from tinygrad import Tensor
+a = Tensor([1,2,3,4]).shard(['CPU', 'CPU:1'], axis=0)
+print(a.numpy())  # should be [1,2,3,4], gets [0,0,0,0]
+"
+```
+
+## Lint/typecheck
+```bash
+python -m mypy tinygrad/
+python -m ruff check .
+```

test/test_tiny.py

@@ -39,14 +39,17 @@ class TestTiny(unittest.TestCase):
     out = Tensor.ones(N).contiguous().sum()
     self.assertEqual(out.item(), N)
 
-  def test_gemm(self, N=getenv("GEMM_N", 64)):
-    a = Tensor.ones(N,N).contiguous()
-    b = Tensor.eye(N).clone()
+  def test_gemm(self, N=getenv("GEMM_N", 64), dtype=dtypes.float):
+    a = Tensor.ones(N,N, dtype=dtype).contiguous()
+    b = Tensor.eye(N, dtype=dtype).clone()
     lst = (out:=a@b).tolist()
     for y in range(N):
       for x in range(N):
         self.assertEqual(lst[y][x], 1.0, msg=f"mismatch at ({y},{x})")
-    self.assertEqual(out.dtype, dtypes.float)
+    self.assertEqual(out.dtype, dtype)
+
+  @unittest.skipIf(Device.DEFAULT == "DSP", "half is broken on DSP")
+  def test_hgemm(self): self.test_gemm(dtype=dtypes.half)
 
   def test_gemv(self, N=getenv("GEMV_N", 64), out_dtype=dtypes.float):
     a = Tensor.ones(1,N).contiguous()

test/unit/test_dtype_weak.py

@@ -11,6 +11,7 @@ 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)):
@@ -22,6 +23,17 @@ 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,5 +83,10 @@ 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,9 +79,10 @@ 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.tag != 1: return None
+  if u.arg[4] is None: 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]), tag=None)
+  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))
   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, tag=1, tc_upcast_axes=tc_upcast_axes)
+                              tc.dims, self.ren.target.device, tc.threads, 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,6 +163,8 @@ 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/mixin/op.py

@@ -284,8 +284,8 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
     X = self.shrink(tuple((-smin(pB,0),smin(pA+s,s)) for (pB,pA),s in zip(pX, self.shape))) if has_neg else self
     pads = tuple((smax(pB,0), smax(pA,0)) for pB,pA in pX) if has_neg else pX
     base = MovementMixin.pad(X, pads)
-    if value == 0: return base
-    if value is not Invalid: base = base.cast(least_upper_dtype(base.dtype, dtypes.from_py(value)))
+    if value is Invalid: return base
+    if value != 0: base = base.cast(least_upper_dtype(base.dtype, dtypes.from_py(value)))
     return MovementMixin.pad(X.const_like(1).cast(dtypes.bool), pads).where(base, base.const_like(value))
 
   def _pad_circular(self, pX:tuple[tuple[sint, sint], ...]) -> Self:

tinygrad/mixin/rand.py

@@ -97,9 +97,10 @@ 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(to_dtype(dtype or self.dtype))
+    return src[0].mul(2*math.pi).cos().mul((1 - src[1]).log().mul(-2).sqrt()).cast(dt)
 
   @classmethod
   def randn(cls, *shape, dtype:DTypeLike|None=None, **kwargs) -> Self:

tinygrad/renderer/cstyle.py

@@ -101,10 +101,11 @@ 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_axes)
+# (name, dims, dtype_in, dtype_out, device, threads, upcast_sizes)
 def wmma_args(uops:list[UOp]):
-  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)
+  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)
 
 class CStyleLanguage(Renderer):
   kernel_typedef: str = "void"
@@ -442,8 +443,7 @@ 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_axes in wmma_args(uops):
-      upcast_sizes = [prod(size for _, size in upcast) for upcast in upcast_axes]
+    for name, (N, M, K), dtype_in, dtype_out, _, _, upcast_sizes in wmma_args(uops):
       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], (*x.arg[4][:2], ((0, x.max_numel()*2),)))).index(i*2)
+          arg=(*x.arg[:4], None)).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,6 +165,10 @@ 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/schedule/indexing.py

@@ -1,7 +1,7 @@
 from typing import Iterator
 import functools, itertools
 from dataclasses import dataclass, field, replace
-from tinygrad.dtype import dtypes, AddrSpace
+from tinygrad.dtype import dtypes, AddrSpace, Invalid
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, graph_rewrite, sint, AxisType, profile_matches
 from tinygrad.uop.ops import consumer_map_from_toposort, gate_kernel_sink
 from tinygrad.uop.symbolic import symbolic, pm_simplify_valid, pm_drop_and_clauses
@@ -89,7 +89,7 @@ def convert_pad_to_where_to_keep_behavior_local(ctx:IndexingContext, x:UOp):
   if x not in ctx.range_map: return None
   bx = create_bufferize_and_index_based_on_ranges(ctx, x)
   valid: UOp = UOp.const(dtypes.bool, True).uprod([r.get_valid() for r in ctx.range_map[x][0]])
-  return valid.where(bx.src[0], UOp.const(x.dtype, 0))
+  return valid.where(bx.src[0], UOp.const(x.dtype, Invalid))
 
 def convert_reduce_to_reduce_with_ranges(ctx:IndexingContext, x:UOp):
   if x.arg[1] == 0: return None

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, _from_np_dtype, _to_np_dtype, PyConst
+from tinygrad.dtype import DType, DTypeLike, dtypes, ConstType, to_dtype, strong_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(self.dtype).contiguous()
+    x = self.cast(strong_dtype(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,10 +255,11 @@ 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=}"
-    fmt = self.dtype.fmt
-    assert fmt is not None, f"no fmt dtype for {self.dtype}"
+    buf = self._buffer()
+    fmt = buf.dtype.fmt
+    assert fmt is not None, f"no fmt dtype for {buf.dtype}"
     assert fmt != "e" or sys.version_info >= (3, 12)
-    return self._data().cast(fmt, self.shape)  # type: ignore[arg-type,return-value]
+    return buf.as_memoryview().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]:
@@ -277,6 +278,10 @@ 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, Invalid, AddrSpace
+from tinygrad.dtype import ConstType, dtypes, DType, DTypeLike, truncate, least_upper_dtype, least_upper_float, strong_dtype, 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,9 +367,8 @@ 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 + (prod([x for _,x in out0]),)
+        return wmma_b + (self.src[2].shape[-1],)
 
       # passthrough ops
       case Ops.MSTACK | Ops.MSELECT | Ops.DETACH | Ops.CONTIGUOUS | Ops.CONTIGUOUS_BACKWARD | Ops.AFTER | Ops.LOAD | \
@@ -601,11 +600,9 @@ 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, 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)))
+  def wmma(a:UOp, b:UOp, acc:UOp, dims:tuple[int, int, int], device:str, threads:int, tc_upcast_axes=None):
     # 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), tag=tag)
+    return UOp(Ops.WMMA, src=(a, b, acc), arg=(dims, a.dtype, device, threads, tc_upcast_axes))
   def _rop(self, op:Ops, axis:tuple[int, ...]):
     # NOTE: we don't allow reduce on 1s axis
     axis = tuple(sorted(axis))
@@ -788,9 +785,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(device=device)
+    ret = self.empty_like(dtype=strong_dtype(self.dtype), 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))
+    return ret.after(ret.store(src.cast(ret.dtype)))
   @recursive_property
   def device(self) -> str|tuple[str, ...]|None:
     if self.op is Ops.PARAM: return self.arg.device

tinygrad/uop/symbolic.py

@@ -85,12 +85,19 @@ pm_data_invalid = PatternMatcher([
   (UPat(GroupOp.Binary, src=(UPat.var("y"), invalid_gate), name="alu"), lambda cond,x,y,alu,i: cond.where(y.alu(alu.op,x), i.cast(alu.dtype))),
   (UPat(GroupOp.Binary-GroupOp.Comparison, src=[invalid_pat, UPat()]), lambda i: i),
   # an Invalid condition poisons the whole where; a gated Invalid condition lifts the gate out
+  # when the gate's valid value is a const (e.g. a pad mask: where(valid, True, Invalid)),
+  # use the else value in don't-care positions so masks work
   (invalid_pat.where(UPat.var("a"), UPat()), lambda i,a: i.cast(a.dtype)),
+  (UPat.var("cond").where(UPat.cvar("x"), invalid_pat).where(UPat.var("a"), UPat.cvar("b")),
+   lambda cond,x,i,a,b: cond.where(x.where(a,b), b)),
   (invalid_gate.where(UPat.var("a"), UPat.var("b")), lambda cond,x,i,a,b: cond.where(x.where(a,b), i.cast(a.dtype))),
   # normalize where(cond, Invalid, val) -> where(~cond, val, Invalid)
   (UPat.var("cond").where(invalid_pat, UPat.var("val")), lambda cond, i, val: cond.logical_not().where(val, i) if val.arg != Invalid else i),
   # lift Invalid out: a.where(cond.where(x, Invalid), c) -> (~a|cond).where(a.where(x, c), Invalid)
   # when a is cond, ~a|cond is True and would drop the Invalid gate (losing the valid), so keep cond as the gate
+  # when c is a const and the gate's valid value is a const (pad mask), use c in don't-care positions
+  (UPat.var("a").where(UPat.var("cond").where(UPat.cvar("x"), invalid_pat), UPat.cvar("c")),
+   lambda cond,i,x,a,c: (cond if a is cond else (a.logical_not()|cond)).where(a.where(x,c), c) if c.arg != Invalid else None),
   (UPat.var("a").where(invalid_gate, UPat.var("c")), lambda cond,i,x,a,c:
    (cond if a is cond else (a.logical_not()|cond)).where(a.where(x,c), i) if c.arg != Invalid else None),
   (UPat.var("a").where(UPat.var("b"), invalid_gate), lambda cond,i,x,a,b: (a|cond).where(a.where(b, x), i) if b.arg != Invalid else None),
@@ -347,7 +354,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 sum(-1 if (res:=parse_valid(v)) is not None and res[0] in other.toposort() else 0 for other in valids)
+  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)
 
 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