From 9506b78d730386cd0661eb6e044eb86ecc241941 Mon Sep 17 00:00:00 2001
From: George Hotz <72895+geohot@users.noreply.github.com>
Date: Fri, 29 May 2026 12:58:05 -0700
Subject: [PATCH] fix viz addrspace (#16437)

* fix viz addrspace

* revert that
---
 tinygrad/codegen/gpudims.py | 2 +-
 tinygrad/uop/ops.py         | 4 ++--
 2 files changed, 3 insertions(+), 3 deletions(-)

diff --git a/tinygrad/codegen/gpudims.py b/tinygrad/codegen/gpudims.py
index 0077e7a570..7ffd44c2cc 100644
--- a/tinygrad/codegen/gpudims.py
+++ b/tinygrad/codegen/gpudims.py
@@ -91,7 +91,7 @@ def add_gpudims(ctx:Renderer, s:UOp):
   subs = {}
   for r in s_topo:
     # look for local INDEXes that are not used in the GLOBAL store, then add them as an INVALID
-    if r.op is Ops.STORE and (idx := r.src[0]).src[0].buf_uop.addrspace == AddrSpace.GLOBAL:
+    if r.op is Ops.STORE and (idx := r.src[0]).src[0].addrspace == AddrSpace.GLOBAL:
       missing_locals = [all_ranges[rng] for rng in local_dims if all_ranges[rng] not in idx.ranges]
       if len(missing_locals):
         assert len(idx.src) == 2, "index has 2 sources"
diff --git a/tinygrad/uop/ops.py b/tinygrad/uop/ops.py
index 98463d3886..4b9ed82f6c 100644
--- a/tinygrad/uop/ops.py
+++ b/tinygrad/uop/ops.py
@@ -764,8 +764,8 @@ class UOp(OpMixin, metaclass=UOpMetaClass):
     if self.op in GroupOp.Movement: return self.src[0].addrspace
     if self.op is Ops.STACK or self.op in GroupOp.Elementwise:
       ad = [x.addrspace for x in self.src if x.addrspace is not None]
-      assert all_same(ad), f"addrspace mismatch in {self.op} -- {ad}"
-      return ad[0] if len(ad) else None
+      if not len(ad) or not all_same(ad): return None
+      return ad[0]
     return None
   @property
   def buf_uop(self) -> UOp: