pad invalid work (glm)

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/null/test_llm_server.py

@@ -7,9 +7,12 @@ class TestLLMServer(unittest.TestCase):
   @classmethod
   def setUpClass(cls):
     cls.mock_tok = Mock()
+    cls.mock_tok.role = Mock(return_value=[100, 101])
     cls.mock_tok.encode = Mock(return_value=[200, 201, 202])
     cls.mock_tok.decode = Mock(return_value="Hello")
     cls.mock_tok.stream_decoder = Mock(return_value=lambda tid=None: "Hello" if tid is not None else "")
+    cls.mock_tok.end_turn = Mock(return_value=[998])
+    cls.mock_tok.prefix = Mock(return_value=[1])
     cls.mock_tok.preset = "llama3"
     cls.mock_tok.bos_id = 1
     cls.mock_tok.eos_id = 999
@@ -20,9 +23,9 @@ class TestLLMServer(unittest.TestCase):
     cls.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 301, 999]))
     cls.mock_model.get_start_pos = Mock(return_value=0)
 
-    from tinygrad.llm.cli import LLMServer, FallbackTemplate
+    from tinygrad.llm.cli import LLMServer
 
-    cls.server = LLMServer(('127.0.0.1', 0), cls.mock_model, "test-model", cls.mock_tok, FallbackTemplate(cls.mock_tok))
+    cls.server = LLMServer(('127.0.0.1', 0), cls.mock_model, "test-model", cls.mock_tok)
     cls.port = cls.server.server_address[1]
     cls.server_thread = threading.Thread(target=cls.server.serve_forever, daemon=True)
     cls.server_thread.start()
@@ -146,6 +149,50 @@ class TestLLMServer(unittest.TestCase):
     self.assertEqual(resp.choices[0].finish_reason, "length")
     self.assertEqual(resp.usage.completion_tokens, 2)
 
+  def test_assistant_prefill(self):
+    """Last assistant message should be treated as prefill (not a completed turn)."""
+    self.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 999]))
+    captured_ids = []
+    orig_generate = self.mock_model.generate.side_effect
+    def capture_generate(ids, **kwargs):
+      captured_ids.extend(ids)
+      return orig_generate(ids, **kwargs)
+    self.mock_model.generate = Mock(side_effect=capture_generate)
+
+    resp = self.client.chat.completions.create(
+      model="test", messages=[
+        {"role": "user", "content": "Hello"},
+        {"role": "assistant", "content": "Sure"}
+      ], stream=False
+    )
+    # prefill tokens should be in ids: role("assistant") + encode("Sure") but NO end_turn after it
+    # and NO extra role("assistant") appended
+    role_tokens = self.mock_tok.role.call_args_list
+    # last role() call should be for "assistant" (the prefill message), not an extra one
+    self.assertEqual(role_tokens[-1], unittest.mock.call("assistant"))
+    # end_turn should be called once less than role() — the prefill assistant msg doesn't get end_turn
+    # NOTE: this is flaky in random order
+    #self.assertEqual(self.mock_tok.end_turn.call_count, self.mock_tok.role.call_count - 1)
+    self.assertIsNotNone(resp.choices[0].message.content)
+
+  def test_assistant_prefill_not_last(self):
+    """Assistant message that's NOT last should be a normal completed turn."""
+    self.mock_model.generate = Mock(side_effect=lambda ids, **kwargs: iter([300, 999]))
+    self.mock_tok.role.reset_mock()
+    self.mock_tok.end_turn.reset_mock()
+    self.client.chat.completions.create(
+      model="test", messages=[
+        {"role": "user", "content": "Hello"},
+        {"role": "assistant", "content": "Sure"},
+        {"role": "user", "content": "Continue"}
+      ], stream=False
+    )
+    # all messages get end_turn, plus an extra role("assistant") at the end
+    # roles: user, assistant, user, assistant(generation prompt) = 4 role calls
+    # end_turns: user, assistant, user = 3 end_turn calls (one per message)
+    self.assertEqual(self.mock_tok.end_turn.call_count, 3)
+    self.assertEqual(self.mock_tok.role.call_count, 4)
+
   def test_models_endpoint(self):
     import requests as req
     resp = req.get(f"http://127.0.0.1:{self.port}/v1/models")

test/null/test_llm_tokenizer.py

@@ -1,5 +1,5 @@
 import unittest, base64, functools, sys
-from tinygrad.llm.cli import SimpleTokenizer, FallbackTemplate
+from tinygrad.llm.cli import SimpleTokenizer
 from tinygrad.helpers import fetch
 
 @unittest.skipIf(sys.platform == 'win32', "fetch race condition on Windows")
@@ -54,11 +54,10 @@ class TestLLMTokenizer(unittest.TestCase):
       "tokenizer.ggml.eos_token_id": 2,
     }
     tok = SimpleTokenizer.from_gguf_kv(kv)
-    template = FallbackTemplate(tok)
-    self.assertEqual(template.role("user"), "[INST]")
+    self.assertEqual(tok.role("user"), [3])
     self.assertEqual(tok.encode("hello"), [5])
-    self.assertEqual(template.end_turn(), "[/INST]")
-    self.assertEqual(template.role("assistant"), "")
+    self.assertEqual(tok.end_turn(), [4])
+    self.assertEqual(tok.role("assistant"), [])
 
   def test_stream_decoder(self):
     """stream_decoder buffers incomplete UTF-8: token 25677 has 3/4 of emoji, token 138 completes it."""

tinygrad/llm/cli.py

@@ -1,13 +1,10 @@
 from __future__ import annotations
 import sys, argparse, codecs, typing, re, unicodedata, json, uuid, time, pathlib
-from typing import TYPE_CHECKING
 from tinygrad import nn
 from tinygrad.uop.ops import UOp, Ops
 from tinygrad.helpers import partition, DEBUG, Timing, GlobalCounters, stderr_log, colored, Context, fetch, profile_marker, getenv
 from tinygrad.viz.serve import TCPServerWithReuse, HTTPRequestHandler
 from tinygrad.llm.model import Transformer
-if TYPE_CHECKING:
-  import jinja2
 
 class SimpleTokenizer:
   def __init__(self, normal_tokens:dict[str, int], special_tokens:dict[str, int], preset:str="llama3",
@@ -67,6 +64,25 @@ class SimpleTokenizer:
     dec = codecs.getincrementaldecoder('utf-8')('replace')
     def _decode(tid:int|None=None) -> str: return dec.decode(self._tok2bytes[tid]) if tid is not None else dec.decode(b'', final=True)
     return _decode
+  def role(self, role:str):
+    if self.preset == 'olmo': return self.encode("<|" + role + "|>\n")  # OLMoE Instruct format
+    if self.preset == 'kimi-k2': return self.encode("<|im_" + role + "|>" + role + "<|im_middle|>")
+    if self.preset == 'qwen2': return self.encode("<|im_start|>" + role + "\n")
+    if self.preset == 'glm4': return self.encode("<|" + role + "|>")
+    if self.preset == 'tekken':
+      if role == 'user': return self.encode("[INST]")
+      if role == 'assistant': return []
+      raise ValueError(f"Unsupported role '{role}' for tokenizer preset '{self.preset}'")
+    return self.encode("<|start_header_id|>" + role + "<|end_header_id|>\n\n")
+  def end_turn(self):
+    if self.preset == 'olmo': return self.encode("\n")
+    if self.preset == 'kimi-k2': return [self.eos_id]
+    if self.preset == 'qwen2': return [self.eos_id] + self.encode("\n")
+    if self.preset == 'glm4': return []
+    if self.preset == 'tekken': return self.encode("[/INST]")
+    return [self.eos_id]
+  def prefix(self) -> list[int]:
+    return ([] if self.bos_id is None else [self.bos_id]) + (self.encode("<sop>") if self.preset == 'glm4' else [])
   def is_end(self, token_id:int) -> bool: return token_id in (self.eos_id, self.eot_id)
 
 models = {
@@ -91,40 +107,6 @@ models = {
 
 # *** simple OpenAI API compatible server with web interface on http://localhost:8000/ ***
 
-class FallbackTemplate:
-  # minimal jinja2.Template-compatible chat template without jinja2, no tool calling support
-  def __init__(self, tok:SimpleTokenizer): self.tok = tok
-  def role(self, role:str) -> str:
-    if self.tok.preset == 'olmo': return "<|" + role + "|>\n"  # OLMoE Instruct format
-    if self.tok.preset == 'kimi-k2': return "<|im_" + role + "|>" + role + "<|im_middle|>"
-    if self.tok.preset == 'qwen2': return "<|im_start|>" + role + "\n"
-    if self.tok.preset == 'glm4': return "<|" + role + "|>"
-    if self.tok.preset == 'tekken':
-      if role == 'user': return "[INST]"
-      if role == 'assistant': return ""
-      raise ValueError(f"Unsupported role '{role}' for tokenizer preset '{self.tok.preset}'")
-    return "<|start_header_id|>" + role + "<|end_header_id|>\n\n"
-  def end_turn(self) -> str:
-    if self.tok.preset == 'olmo': return "\n"
-    if self.tok.preset == 'kimi-k2': return self.tok.decode([self.tok.eos_id])
-    if self.tok.preset == 'qwen2': return self.tok.decode([self.tok.eos_id]) + "\n"
-    if self.tok.preset == 'glm4': return ""
-    if self.tok.preset == 'tekken': return "[/INST]"
-    return self.tok.decode([self.tok.eos_id])
-  def render(self, messages:list[dict], tools=None, add_generation_prompt:bool=True) -> str:
-    out = self.tok.decode([] if self.tok.bos_id is None else [self.tok.bos_id]) + ("<sop>" if self.tok.preset == 'glm4' else "")
-    for msg in messages:
-      out += self.role(msg["role"])
-      content = msg.get("content")
-      if isinstance(content, str): out += content
-      elif isinstance(content, list):
-        for c in content:
-          if c["type"] == "text": out += c["text"]
-          else: raise RuntimeError(f"unhandled type: {c['type']}")
-      elif content is not None: raise RuntimeError(f"unknown content type: {type(content)}")
-      out += self.end_turn()
-    return out + self.role("assistant") if add_generation_prompt else out
-
 class Handler(HTTPRequestHandler):
   server: LLMServer
   def log_request(self, code='-', size='-'): pass
@@ -159,13 +141,25 @@ class Handler(HTTPRequestHandler):
                f"out:{len(out):5d}  {colored('--', 'BLACK')}  total:{et-st:6.2f}s\n")
 
   def do_POST(self):
+    tok = self.server.tok
     raw_body = self.rfile.read(int(self.headers.get("Content-Length", "0")))
     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":
-      # render and tokenize
-      rendered = self.server.template.render(messages=body["messages"], tools=body.get("tools"), add_generation_prompt=True)
-      ids: list[int] = self.server.tok.encode(rendered)
+      # 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")
@@ -183,8 +177,8 @@ class Handler(HTTPRequestHandler):
       raise RuntimeError(f"unhandled path {self.path}")
 
 class LLMServer(TCPServerWithReuse):
-  def __init__(self, server_address:tuple, model:Transformer, model_name:str, tok:SimpleTokenizer, template:jinja2.Template|FallbackTemplate):
-    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():
@@ -200,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)
 
-  # use the model's chat template if jinja2 is available (enables model-specific formatting)
-  template: jinja2.Template|FallbackTemplate = FallbackTemplate(tok)
-  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, **kwargs)  # 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)
-      env.globals['bos_token'] = tok.decode([tok.bos_id]) if tok.bos_id is not None else ""
-      env.globals['eos_token'] = tok.decode([tok.eos_id])
-      template = env.from_string(ct)
-    except ImportError: print("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
@@ -227,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:
@@ -245,19 +224,16 @@ def main():
     exit(0)
 
   # interactive chat
-  messages: list[dict] = []
+  ids: list[int] = tok.prefix()
   while 1:
-    try: messages.append({"role":"user", "content":input('>>> ')})
-    except EOFError: break
-    ids = tok.encode(template.render(messages=messages, add_generation_prompt=True))
-    reply, dec = "", tok.stream_decoder()
+    try:
+      ids += tok.role("user") + tok.encode(input('>>> ')) + tok.end_turn() + tok.role("assistant")
+    except EOFError:
+      break
+    dec = tok.stream_decoder()
     for next_id in model.generate(ids):
-      if tok.is_end(next_id):
-        sys.stdout.write(dec() + "\n\n")
-        break
-      reply += (piece := dec(next_id))
-      sys.stdout.write(piece)
+      sys.stdout.write(dec(next_id) if not tok.is_end(next_id) else dec() + "\n\n")
       sys.stdout.flush()
-    messages.append({"role":"assistant", "content":reply})
+      if tok.is_end(next_id): break
 
 if __name__ == "__main__": main()

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/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/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),