low rank call in kernel graph

fix custom_kernel wrapper
call has rank <= 1
2026-08-28 06:16:08 +00:00 · 2026-08-27 10:34:54 -07:00 · 2026-08-27 10:18:37 -07:00 · 2026-08-27 10:15:35 -07:00
58 changed files with 326 additions and 584 deletions
@@ -194,29 +194,21 @@ runs:
        echo "pkgs=$pkgs" >> "$GITHUB_OUTPUT"
        echo "hash=$(echo -n "$pkgs" | sha256sum | cut -d' ' -f1)" >> "$GITHUB_OUTPUT"

-        installed=true
-        for pkg in $pkgs; do
-          info=$(dpkg-query -W -f='${db:Status-Abbrev} ${Version}' "$pkg" 2> /dev/null || true)
-          echo "${pkg}: ${info:-not in dpkg database}"
-          [[ "$info" == ii* ]] || installed=false
-        done
-        echo "installed=$installed" >> "$GITHUB_OUTPUT"
-
    - name: Cache apt (PR)
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && github.event_name == 'pull_request' && steps.apt-pkgs.outputs.installed == 'false'
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && github.event_name == 'pull_request'
      uses: actions/cache/restore@v5
      with:
        path: /var/cache/apt/archives/
        key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}
    - name: Cache apt
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && github.event_name != 'pull_request' && steps.apt-pkgs.outputs.installed == 'false'
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && github.event_name != 'pull_request'
      uses: actions/cache@v5
      with:
        path: /var/cache/apt/archives/
        key: ${{ runner.os }}-${{ runner.arch }}-apt-${{ steps.apt-pkgs.outputs.hash }}-${{ env.CACHE_VERSION }}

    - name: Run apt Update + Install
-      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true') && steps.apt-pkgs.outputs.installed == 'false'
+      if: runner.os == 'Linux' && (inputs.opencl == 'true' || inputs.amd == 'true' || inputs.webgpu == 'true' || inputs.llvm == 'true' || inputs.qemu == 'true' || inputs.ninja == 'true')
      shell: bash
      run: |
        sudo apt -qq update || true
@@ -40,10 +40,10 @@ jobs:
      run: sudo apt-get install -y --no-install-recommends libclang-20-dev llvm-20-dev hip-dev libusb-1.0-0-dev libdrm-dev liburing-dev
    - name: Regenerate autogen files
      run: |
-        find tinygrad/runtime/autogen -type f -name "*.py" -not -path "*/amd/*" -not -name "__init__.py" -not -name "metal.py" -not -name "iokit.py" -not -name "corefoundation.py" -not -name "libclang.py" -delete
+        find tinygrad/runtime/autogen -type f -name "*.py" -not -path "*/amd/*" -not -name "__init__.py" -not -name "comgr.py" -not -name "metal.py" -not -name "iokit.py" -not -name "corefoundation.py" -not -name "libclang.py" -delete
        python3 -c "from tinygrad.runtime.autogen import opencl"
        python3 -c "from tinygrad.runtime.autogen import cuda, nvrtc, nvjitlink, nv_570, nv_580, nv_610, nv"
-        python3 -c "from tinygrad.runtime.autogen import comgr, comgr_3, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd, amdgpu_drm"
+        python3 -c "from tinygrad.runtime.autogen import comgr_3, hsa, hip, amd_gpu, sqtt, rocprof, amdgpu_kd, amdgpu_drm"
        python3 -c "from tinygrad.runtime.autogen.am import *"
        python3 -c "from tinygrad.runtime.autogen.nv_regs import *"
        python3 -c "from tinygrad.runtime.autogen import libc, kfd, io_uring, pci, vfio"
@@ -102,3 +102,42 @@ jobs:
      with:
        name: autogen-macos-patch
        path: autogen-macos.patch
+
+  autogen-comgr-2:
+    name: In-tree Autogen (comgr 2)
+    runs-on: ubuntu-24.04
+    timeout-minutes: 15
+    steps:
+    - name: Checkout Code
+      uses: actions/checkout@v6
+    - name: Setup Environment
+      uses: ./.github/actions/setup-tinygrad
+      with:
+        key: 'autogen-comgr'
+    - name: Install autogen support packages
+      run: |
+        wget https://repo.radeon.com/rocm/rocm.gpg.key -O - | gpg --dearmor | sudo tee /etc/apt/keyrings/rocm.gpg > /dev/null
+        sudo tee /etc/apt/sources.list.d/rocm.list <<EOF
+        deb [arch=amd64 signed-by=/etc/apt/keyrings/rocm.gpg] https://repo.radeon.com/rocm/apt/6.2 $(lsb_release -cs) main
+        EOF
+        echo -e 'Package: *\nPin: release o=repo.radeon.com\nPin-Priority: 600' | sudo tee /etc/apt/preferences.d/rocm-pin-600
+        sudo apt -qq update || true
+        sudo apt-get install -y --no-install-recommends libclang-20-dev comgr
+    - name: Regenerate autogen files
+      run: |
+        rm tinygrad/runtime/autogen/comgr.py
+        python3 -c "from tinygrad.runtime.autogen import comgr"
+    - name: Check for differences
+      run: |
+        if ! git diff --quiet; then
+          git diff
+          git diff > autogen-comgr2.patch
+          echo "Autogen mismatch detected. Patch available at: ${{ github.server_url }}/${{ github.repository }}/actions/runs/${{ github.run_id }}#artifacts"
+          exit 1
+        fi
+    - name: Upload patch artifact
+      if: failure()
+      uses: actions/upload-artifact@v7
+      with:
+        name: autogen-comgr2-patch
+        path: autogen-comgr2.patch
@@ -94,7 +94,7 @@ jobs:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: '0'
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -140,7 +140,7 @@ jobs:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: '0'
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -188,7 +188,7 @@ jobs:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: '0'
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -230,7 +230,7 @@ jobs:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: '0'
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -275,7 +275,7 @@ jobs:
        shell: bash -e -o pipefail {0}
    env:
      DEV: ${{ matrix.dev }}
-      HCQ2: '0'
+      HCQ2: ${{ matrix.dev == 'AMD' && '1' || '0' }}
    if: github.repository_owner == 'tinygrad'
    steps:
    - name: Checkout Code
@@ -390,10 +390,10 @@ jobs:
        run: |
          parallel --link --tagstring '[{1}]' '{2}' \
            ::: llama 'llama q4' qwen3.5 qwen \
-            ::: $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model llama3.2:1b | tee /dev/stderr | grep -i rooster' \
-                $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model llama3.2:1b-q4 | tee /dev/stderr | grep -i rooster' \
-                $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model qwen3.5:0.8b | tee /dev/stderr | grep -i rooster' \
-                $'echo "What\'s a female chicken called? Answer with only one word." | python3 -m tinygrad.llm --no_chat_template --model qwen3:0.6b | tee /dev/stderr | grep -i hen'
+            ::: $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --model llama3.2:1b | tee /dev/stderr | grep -i rooster' \
+                $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --model llama3.2:1b-q4 | tee /dev/stderr | grep -i rooster' \
+                $'echo "What\'s a male chicken called? Answer with only one word." | python3 -m tinygrad.llm --model qwen3.5:0.8b | tee /dev/stderr | grep -i rooster' \
+                $'echo "What\'s a female chicken called? Answer with only one word." | python3 -m tinygrad.llm --model qwen3:0.6b | tee /dev/stderr | grep -i hen'
                # NOTE: qwen is dumb and only knows about female chickens

 # ****** Models Tests ******
@@ -2,7 +2,7 @@ import os, pytest, signal, threading

@pytest.hookimpl(wrapper=True)
 def pytest_runtest_call(item):
-  t = threading.Timer(int(os.getenv("TEST_TIMEOUT", 90)), os.kill, args=(os.getpid(), signal.SIGABRT))
+  t = threading.Timer(int(os.getenv("TEST_TIMEOUT", 300)), os.kill, args=(os.getpid(), signal.SIGABRT))
  t.start()
  try: yield
  finally:
@@ -1,6 +1,8 @@
 from __future__ import annotations
 import functools, pathlib
-from tinygrad import Tensor
+from dataclasses import replace
+from tinygrad import Tensor, dtypes
+from tinygrad.uop.ops import shape_to_shape_arg
 from tinygrad.runtime.support.compiler_amd import HIPCCCompiler

 FP8_MAX = 448.0
@@ -10,7 +12,7 @@ NUM_WG, THREADS_PER_WG = 1024, 256
@functools.cache
 def _local_abs_max_fxn(x_p, device):
  x = Tensor(x_p, device=device)
-  inner = Tensor(x.uop.src[0]) if x.uop.axis is not None else x  # the per-shard view of the flat param
+  inner = Tensor(x.uop.replace(src=(shape_to_shape_arg(x.uop.shard_shape),), arg=replace(x.uop.arg, axis=None))) if x.uop.axis is not None else x
  return (inner.abs().max(),)

 def local_abs_max(x:Tensor) -> Tensor:
@@ -1,12 +1,10 @@
 from tinygrad import Tensor
-from tinygrad.helpers import NUM_CPU_THREADS
 from tinygrad.tensor import _to_np_dtype
 from tinygrad.nn.onnx import OnnxRunner, OnnxValue
 import numpy as np
 import onnxruntime as ort
 ort_options = ort.SessionOptions()
 ort_options.log_severity_level = 3
-ort_options.intra_op_num_threads = NUM_CPU_THREADS.value

 def get_example_inputs(graph_inputs:dict[str, OnnxValue], config={}):
  """
@@ -50,10 +50,10 @@ All nodes in the tinygrad graph are \textbf{UOps}. A UOp is a tuple $(\mathrm{op
 \toprule
 \textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
 \midrule
-\op{Param} & () & slot, dtype, size?, device?, addrspace? &
-  Placeholder with flat storage of $\mathrm{size}$ elements. Substituted in \op{Call}. \\[4pt]
-\op{Buffer} & () & slot, dtype, size, device, addrspace &
-  Concrete buffer slot with flat storage of $\mathrm{size}$ elements. \\
+\op{Param} & $(\mathbf{s})$ & slot, dtype, device?, addrspace? &
+  Placeholder with shape $\mathbf{s}$. Substituted in \op{Call}. \\[4pt]
+\op{Buffer} & $(\mathbf{s})$ & slot, dtype, device, addrspace &
+  Concrete buffer slot with shape $\mathbf{s}$. If device is a tuple, it creates the fully sized buffer across multiple devices. \\
 \op{Const} & () & value, dtype &
  A scalar constant with shape $(\ )$. \\
  & & & Form vector consts with \op{Stack} \\
@@ -102,9 +102,8 @@ All nodes in the tinygrad graph are \textbf{UOps}. A UOp is a tuple $(\mathrm{op
 \toprule
 \textbf{Op} & \textbf{src} & \textbf{arg} & \textbf{Semantics} \\
 \midrule
-\op{Function} & (body, $a_0$, $a_1$, \ldots) & --- & Substitute each \op{Param} $k$ in \op{Tuple} body with $a_k$. Gradient-able. \\
-\op{Call} & (body, $a_0$, $a_1$, \ldots) & --- & Opaque invocation of a compiled kernel or custom function. \\
-\op{Tuple} & $(v_0, v_1, \ldots)$ & --- & Pack values; required as \op{Function} body to return a value. \\
+\op{Call} & (body, $a_0$, $a_1$, \ldots) & --- & Substitute each \op{Param} $k$ in body with $a_k$. \\
+\op{Tuple} & $(v_0, v_1, \ldots)$ & --- & Pack values; required as \op{Call} body to return a value. \\
 \op{GetTuple} & $(T,)$ & idx & Extract element at idx from a \op{Tuple}. \\
 \bottomrule
 \end{tabular}
@@ -256,9 +255,9 @@ Every UOp has a \textbf{dtype}, \textbf{shape}, \textbf{device}, \textbf{addrspa
 \toprule
 \textbf{Op} & \textbf{dtype} & \textbf{shape} & \textbf{device} & \textbf{min\_max} \\
 \midrule
-\op{Buffer} & from arg & from arg ($\mathrm{size}$) & from arg & dtype range \\
+\op{Buffer} & from arg & from $\mathrm{src}[0]$ & from arg & dtype range \\
 \op{Const} & from arg & $()$ & \textsc{null} & $[v, v]$ \\
-\op{Param} & from arg & from arg ($\mathrm{size}$) & from arg & from src or dtype range \\[3pt]
+\op{Param} & from arg & from $\mathrm{src}[0]$ & from arg & from src or dtype range \\[3pt]
 Movement ops & $\mathrm{src}[0].\mathrm{dtype}$ & (see op) & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
 \op{Unshard} & $\mathrm{src}[0].\mathrm{dtype}$ & $\mathrm{src}[0]$, each $a_k \times n_k$ & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
 \op{Reduce} & $\mathrm{src}[0].\mathrm{dtype}$ & remove first $n$ axes & $\mathrm{src}[0].\mathrm{device}$ & dtype range \\[3pt]
@@ -272,7 +271,7 @@ ALU unary & $\mathrm{src}[0].\mathrm{dtype}$ & $\mathrm{src}[0].\mathrm{shape}$
 Other binary & $\mathrm{src}[0].\mathrm{dtype}$ & broadcast & $\mathrm{src}[0].\mathrm{device}$ & dtype range \\
 \op{CmpLt}, \op{CmpNe} & bool & broadcast & $\mathrm{src}[0].\mathrm{device}$ & from intervals \\
 \op{Where} & $\mathrm{src}[1].\mathrm{dtype}$ & broadcast & $\mathrm{src}[0].\mathrm{device}$ & $[\min(b,c),\, \max(B,C)]$ \\[3pt]
-\op{Function}, \op{Call} & $\mathrm{src}[0].\mathrm{dtype}$ & substitute \op{Param} shapes & $\mathrm{src}[1].\mathrm{device}$ & dtype range \\
+\op{Call} & $\mathrm{src}[0].\mathrm{dtype}$ & substitute \op{Param} shapes & $\mathrm{src}[1].\mathrm{device}$ & dtype range \\
 \op{Range} & index & $()$ & \textsc{null} & $[0,\, n{-}1]$ \\
 \op{Index} & $\mathrm{src}[0].\mathrm{dtype}$ & remaining dims & $\mathrm{src}[0].\mathrm{device}$ & $\mathrm{src}[0]$ \\
 \op{Store} & void & $()$ & $\mathrm{src}[0].\mathrm{device}$ & --- \\
@@ -421,7 +420,7 @@ def allreduce(T):
 %% ============================================================
 \subsection*{{\color{callblue}The \texttt{@function} Decorator} \normalfont\small--- graph capture via tracing}

-The \texttt{@function} decorator transforms a Python function on Tensors into a single \op{Function} node.
+The \texttt{@function} decorator transforms a Python function on Tensors into a single \op{Call} node.

 \begin{lstlisting}
@function
@@ -436,11 +435,11 @@ When \texttt{f(x, y)} is called, the decorator:
  \item \textbf{Runs the function} lazily (no device execution), building a UOp graph from the result.
  \item \textbf{Parameterizes}: replaces each input UOp with a \op{Param}$(k)$ placeholder.
  \item \textbf{Wraps the body} in a \op{Tuple} (even for single returns) and creates\\
-        \op{Function}(\op{Tuple}(body), $x$, $y$).
+        \op{Call}(\op{Tuple}(body), $x$, $y$).
  \item \textbf{Returns} the result via \op{GetTuple}$(0)$, or one \op{GetTuple} per element for tuple returns.
 \end{enumerate}

-The result is a reusable graph fragment: the body contains only \op{Param} references, not concrete buffers. At schedule time, the \op{Function} is resolved by substituting each \op{Param}$(k)$ back with its corresponding argument $a_k$, or lowered into an opaque \op{Call} if it is to be compiled as a reusable kernel.
+The result is a reusable graph fragment: the body contains only \op{Param} references, not concrete buffers. At schedule time, the \op{Call} is resolved by substituting each \op{Param}$(k)$ back with its corresponding argument $a_k$, or compiled as a reusable kernel.

 %% ============================================================
 \subsection*{Lowering Pipeline \normalfont\small--- from Tensor graph to machine code}
@@ -152,7 +152,7 @@ class TestDSPcodePatterns(unittest.TestCase):

  def test_global_atomic_add_f32_parsing(self):
    """Test GLOBAL_ATOMIC_ADD_F32 keeps memory values in float dtype."""
-    vmem = UOp.param(2, dtypes.uint32, 1024)
+    vmem = UOp.param(2, dtypes.uint32, (1024,))
    srcs = {
      'ADDR': UOp.const(0, dtypes.uint64),
      'DATA': UOp.const(0x3f800000, dtypes.uint32),
@@ -183,7 +183,7 @@ class TestDSPcodePatterns(unittest.TestCase):
  def test_mem_read_parsing(self):
    """Test MEM[addr].type read expression parsing."""
    # Create a mock LDS buffer
-    lds = UOp.param(3, dtypes.uint32, 16384)
+    lds = UOp.param(3, dtypes.uint32, (16384,))
    addr = UOp.const(0, dtypes.uint32)
    vrs = {'_lds': lds, 'ADDR': addr, 'OFFSET': UOp.const(0, dtypes.uint32)}

@@ -218,7 +218,7 @@ class TestDSPcodePatterns(unittest.TestCase):
    pcode = PCODE.get(DSOp.DS_LOAD_2ADDR_B32)
    self.assertIsNotNone(pcode)
    assert pcode is not None
-    lds = UOp.param(3, dtypes.uint32, 16384)
+    lds = UOp.param(3, dtypes.uint32, (16384,))
    srcs = {
      'ADDR': UOp.const(0, dtypes.uint32),
      'OFFSET0': UOp.const(0, dtypes.uint32),
@@ -299,7 +299,7 @@ class TestConcatWidthParsing(unittest.TestCase):
      self.assertIs(parsed.simplify(), UOp.const(expected, dtypes.uint32))

  def test_permlane64_wave64_pcode_indices(self):
-    vgpr = UOp.param(0, dtypes.uint32, 256)
+    vgpr = UOp.param(0, dtypes.uint32, (256,))
    srcs = {
      'SRC0': UOp.const(0, dtypes.uint32),
      'VDST': UOp.const(1, dtypes.uint32),
@@ -330,7 +330,7 @@ class TestAllPcode(unittest.TestCase):
  def _make_srcs(self):
    """Create dummy source variables for pcode parsing."""
    u32, u64 = lambda v=0: UOp.const(v, dtypes.uint32), lambda v=0: UOp.const(v, dtypes.uint64)
-    lds = UOp.param(3, dtypes.uint32, 16384)
+    lds = UOp.param(3, dtypes.uint32, (16384,))
    return {'laneId': u32(), 'laneID': u32(), 'S0': u32(), 'S1': u32(), 'S2': u32(), 'S3': u32(), 'SRC0': u32(),
            'D0': u32(), 'D1': u32(), 'DST': u32(), 'VDST': u32(), 'SDST': u32(),
            'VCC': u64(), 'VCCZ': u32(), 'EXEC': u64(), 'EXEC_LO': u32(), 'EXECZ': u32(), 'SCC': u32(),
@@ -258,11 +258,6 @@ class TestDoubleDType(TestDType):
      a = [2, 3, 4]
      np.testing.assert_allclose(func(Tensor(a, dtype=self.DTYPE)).numpy(), func(torch.tensor(a, dtype=torch.float64)), rtol=1e-12, atol=1e-12)

-  def test_float32_compare_selecting_float64(self):
-    a = Tensor([1.0, 2.0, 5.0, 9.0], dtype=dtypes.float32)
-    p, q = Tensor([10., 20., 30., 40.], dtype=self.DTYPE), Tensor([50., 60., 70., 80.], dtype=self.DTYPE)
-    _test_op(lambda: (a < 3.0).where(p, q), self.DTYPE, [10., 20., 70., 80.])
-
  def test_float64_to_float32_cast_inf(self):
    _test_op(lambda: Tensor([3.4e40, 3.4e38, 1, 0], dtype=dtypes.float64).cast(dtypes.float32),
             dtypes.float32, [float('inf'), 3.4e38, 1, 0])
@@ -428,11 +423,6 @@ class TestDtypeUsage(unittest.TestCase):
        t = Tensor([[1, 2], [3, 4]], dtype=d)
        (t*t).max().item()

-  def test_where_float16_compare_to_const(self):
-    # t > 0 is CMPLT(0, t): the float16 operand is on the right
-    t = Tensor([-1.0, 1.0], dtype=dtypes.float16)
-    np.testing.assert_equal((t > 0).where(Tensor.ones(2, dtype=dtypes.float16), Tensor.zeros(2, dtype=dtypes.float16)).numpy(), [0.0, 1.0])
-
@unittest.skipUnless(dtypes.bfloat16 in supported_dtypes, f"no bfloat16 on {Device.DEFAULT}")
 class TestOpsBFloat16(unittest.TestCase):
  def test_cast(self):
@@ -46,7 +46,7 @@ class TestIselX86(unittest.TestCase):
  # complex address is [base + index*scale + displacement]
  def test_complex_address(self):
    a = UOp.variable("a", 0, 0, dtypes.int32)
-    load = UOp.param(0, dtypes.int32, 16).index(a + UOp.cconst(1, dtypes.int32)).load()
+    load = UOp.param(0, dtypes.int32, (16,)).index(a + UOp.cconst(1, dtypes.int32)).load()
    n = self.isel_rewrite(load)
    # displacement is the constant in "a" scaled to the buffer element size, dtype is int8 when the value fits otherwise int32
    self.assertTrue(n.src[2].dtype is dtypes.int8 and n.src[2].src[0].op is Ops.CONST and n.src[2].src[0].val == 4)
@@ -11,16 +11,16 @@ from tinygrad.codegen import to_program
 class TestLinearizerFailure(unittest.TestCase):
  @unittest.skipUnless(Device.DEFAULT == "METAL", "only tested on METAL")
  def test_failure_beam_mnist(self):
-    c0 = UOp.param(0, dtypes.uchar, 4014080)
+    c0 = UOp.param(0, dtypes.uchar, (4014080,))
    c1 = UOp.range(UOp.const(512), 0, AxisType.GLOBAL)
    c2 = UOp.range(UOp.const(784), 1, AxisType.GLOBAL)
    c3 = UOp.range(UOp.const(10), 3, AxisType.GLOBAL)
-    c4 = UOp.param(1, dtypes.int, 512)
+    c4 = UOp.param(1, dtypes.int, (512,))
    c5 = c4.index(c1.valid(UOp.const(True)))
    c6 = UOp.range(UOp.const(6000), 1004, AxisType.REDUCE)
    c7 = UOp.range(UOp.const(3750), 2006, AxisType.REDUCE)
    c8 = UOp.range(UOp.const(16), 2007, AxisType.GROUP_REDUCE)
-    c9 = UOp.param(2, dtypes.uchar, 47040000)
+    c9 = UOp.param(2, dtypes.uchar, (47040000,))
    c10 = c9.index((((c3*UOp.const(4704000))+c2)+(c6*UOp.const(784))).valid(UOp.const(True)))
    c11 = c5.alu(Ops.CMPNE, ((((c3*UOp.const(6000))+c6)+((c7*UOp.const(16))+c8)).alu(Ops.CMPLT, UOp.const(59999)).where(UOp.const(0).cast(dtypes.int), UOp.const(1).cast(dtypes.int)).reduce(c7, c8, arg=Ops.ADD)+UOp.const(-1).cast(dtypes.int))).where(UOp.const(0).cast(dtypes.uchar), c10).reduce(c6, arg=Ops.ADD)
    c12 = c0.index((((c1*UOp.const(7840))+(c2*UOp.const(10)))+c3).valid(UOp.const(True))).store(c11).end(c1, c2, c3)
@@ -2810,7 +2810,7 @@ class TestOps(unittest.TestCase):
        lambda x: Tensor.interpolate(x, size=out_sz, mode="linear"))

  def test_interpolate_linear_corners_aligned(self):
-    for in_sz, out_sz in [((52,),(29,)), ((29,),(52,)), ((29,),(1,))]:
+    for in_sz, out_sz in [((52,),(29,)), ((29,),(52,))]:
      helper_test_op([(2,3)+in_sz],
        lambda x: torch.nn.functional.interpolate(x, size=out_sz, mode="linear", align_corners=True),
        lambda x: Tensor.interpolate(x, size=out_sz, mode="linear", align_corners=True))
@@ -2963,10 +2963,6 @@ class TestOps(unittest.TestCase):
    data = [math.inf, -math.inf, math.nan]
    helper_test_op((), lambda: torch.tensor(data)[torch.tensor([0, 1, 2])], lambda: Tensor(data)[Tensor([0, 1, 2])])

-  def test_fancy_indexing_index_dtypes(self):
-    helper_test_op((), lambda: torch.tensor([10., 20., 30., 40.])[torch.tensor([1, 2, 3, 0])],
-                       lambda: Tensor([10., 20., 30., 40.])[Tensor([1, 2, 3, 0], dtype=dtypes.uint8)])
-
  @slow_test
  def test_slice_fancy_indexing_no_dim_collapse(self):
    a,b,c,d,e,i,j,k,o,p = self._get_index_randoms()
@@ -23,8 +23,8 @@ def _test_uop_result(inputs:list[Tensor], sink:UOp, local_size=None):

 def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp):
  dtype = alu_src_uops[0].dtype
-  a = UOp.param(0, dtype, 1)
-  b = UOp.param(1, dtype, 1)
+  a = UOp.param(0, dtype, (1,))
+  b = UOp.param(1, dtype, (1,))
  idx = UOp.const(0)
  ld = b.index(idx).load()
  alu = ld.alu(alu_op, *alu_src_uops)
@@ -34,7 +34,7 @@ def _setup_and_test_alu(alu_op:Ops, input_val:ConstType, *alu_src_uops:UOp):
 class TestRendererFailures(unittest.TestCase):
  @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
  def test_gated_store_with_alu(self):
-    a = UOp.param(0, dtypes.int, 4)
+    a = UOp.param(0, dtypes.int, (4,))
    gate_alu = (lidx0:=UOp.special(4, 'lidx0')).ne(0)
    gated_alu_store = UOp(Ops.STORE, src=(a.index(lidx0.valid(gate_alu)), UOp.const(1).cast(dtypes.int)))
    sink = UOp(Ops.SINK, src=(gated_alu_store,), arg=KernelInfo())
@@ -43,7 +43,7 @@ class TestRendererFailures(unittest.TestCase):

  @unittest.skipIf(not isinstance(Device[Device.DEFAULT].renderer, (PTXRenderer, PythonRenderer)), "test is for ptx or python renderer")
  def test_gated_store_with_alu_2d(self):
-    a = UOp.param(0, dtypes.int, 8)
+    a = UOp.param(0, dtypes.int, (8,))
    gate_alu_0 = (lidx0:=UOp.special(4, 'lidx0')).ne(0)
    gate_alu_1 = (lidx1:=UOp.special(2, 'lidx1')).ne(0)
    gated_alu_store = UOp(Ops.STORE, src=(a.index((lidx0+lidx1*4).valid(gate_alu_0&gate_alu_1)), UOp.const(1).cast(dtypes.int)))
@@ -78,7 +78,7 @@ class TestCStyleFailures(unittest.TestCase):
@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, WGSLRenderer), "tests for wgsl renderer")
 class TestWGSLFailures(unittest.TestCase):
  def test_folded_packed_store(self):
-    b = UOp.param(0, dtypes.char, 4)
+    b = UOp.param(0, dtypes.char, (4,))
    idx = b.index(UOp.const(0).cast(dtypes.int))
    store = UOp.store(idx, idx.cast(dtypes.uint32).load() & UOp.const(0xffffff00).cast(dtypes.uint32))
    src = Device[Device.DEFAULT].renderer.render(UOp.sink(store, arg=KernelInfo()).toposort())
@@ -93,9 +93,9 @@ class TestWGSLFailures(unittest.TestCase):

  # WGSL has a specific select(alt, val, gate) ternary operator instead of gate?val:alt
  def test_gated_load(self):
-    a = UOp.param(0, dtypes.int, 4)
-    b = UOp.param(1, dtypes.int, 4)
-    c = UOp.param(2, dtypes.int, 4)
+    a = UOp.param(0, dtypes.int, (4,))
+    b = UOp.param(1, dtypes.int, (4,))
+    c = UOp.param(2, dtypes.int, (4,))
    lidx0 = UOp.special(4, "lidx0")
    gate = lidx0.ne(0)
    alt = c.index(lidx0).load()
@@ -110,7 +110,7 @@ class TestWGSLFailures(unittest.TestCase):
 class TestPTXFailures(unittest.TestCase):
  @unittest.skip("INDEX can only have a gate ALU parent, not an IF")
  def test_gated_store_with_if(self):
-    a = UOp.param(0, dtypes.int, 4)
+    a = UOp.param(0, dtypes.int, (4,))
    gate_alu = (lidx0:=UOp.special(4, 'lidx0')).ne(0)
    val = UOp.const(1).cast(dtypes.int)
    if_uop = UOp(Ops.IF, src=(gate_alu,))
@@ -379,30 +379,27 @@ class TestCopyFolding(unittest.TestCase):
    check_schedule(a.clone(), 1, filter_sink=False)

  def test_shrink_copy(self):
-    a = Tensor.arange(4).clone("CPU:1").realize()
-    b = a.to("CPU:2").shrink(((1, 3),)).to("CPU:3")
-    GlobalCounters.reset()
-    run_linear(*check_schedule(b, 3, filter_sink=False))
-    # extra E kernel, copy exactly 4 bytes
-    self.assertEqual(GlobalCounters.global_mem, 4*4 + 2*4*2 + 2*4)
-    self.assertListEqual(b.tolist(), [1, 2])
+    a = Tensor.arange(4)
+    view = a.shrink(((0, 2),))
+    b = view.clone()
+    run_linear(*check_schedule(b, 1, filter_sink=False))
+    self.assertEqual(b.uop.base.buffer.size, 2)
+    self.assertEqual(b.uop.numel(), 2)
+    self.assertListEqual(b.tolist(), [0, 1])

  def test_expanded_copy(self):
-    a = Tensor.arange(4).clone("CPU:1").realize()
-    b = a.to("CPU:2").reshape(4, 1).expand(4, 2).to("CPU:3")
-    GlobalCounters.reset()
-    run_linear(*check_schedule(b, 3, filter_sink=False))
-    # TODO: expands before copy
-    self.assertEqual(GlobalCounters.global_mem, 4*4 + (4*4 + 8*4) + 8*4)
-    self.assertListEqual(b.tolist(), [[0, 0], [1, 1], [2, 2], [3, 3]])
+    a = Tensor.arange(2)
+    view = a.reshape(2, 1).expand(2, 2)
+    b = view.clone()
+    run_linear(*check_schedule(b, 1, filter_sink=False))
+    self.assertEqual(b.uop.base.buffer.size, 4)
+    self.assertEqual(b.uop.numel(), 4)
+    self.assertListEqual(b.tolist(), [[0, 0], [1, 1]])

  def test_permuted_copy(self):
-    a = Tensor.arange(4).clone("CPU:1").realize()
-    b = a.to("CPU:2").reshape(2, 2).permute(1, 0).to("CPU:3")
-    GlobalCounters.reset()
-    run_linear(*check_schedule(b, 3, filter_sink=False))
-    # permutes before copy
-    self.assertEqual(GlobalCounters.global_mem, 4*4 + (4*4 + 4*4) + 4*4)
+    a = Tensor.arange(4)
+    b = a.reshape(2, 2).permute(1, 0)
+    b.realize()
    self.assertListEqual(b.tolist(), [[0, 2], [1, 3]])

  def test_permute_on_disk(self):
@@ -20,7 +20,7 @@ def run_uops(uops_list:list[UOp], bufs:list[Buffer]):

 def uop(uops:list[UOp], op:Ops, dtype:Optional[DType], src:tuple[UOp, ...], arg:Any=None) -> UOp:
  if op is Ops.CONST: uops.append(UOp.const(arg).cast(dtype))
-  elif op is Ops.PARAM: uops.append(UOp.param(arg, dtype, 1))
+  elif op is Ops.PARAM: uops.append(UOp.param(arg, dtype, shape=(1,)))
  else: uops.append(UOp(op, dtype, tuple(src), arg))
  return uops[-1]

@@ -62,7 +62,7 @@ def _test_uops_result(output_dtype, uops, res):
 class TestBitcastBufferView(unittest.TestCase):
  @Context(SPEC=2)
  def test_render(self):
-    buf = UOp.param(0, dtypes.uint32, 4)
+    buf = UOp.param(0, dtypes.uint32, (4,))
    uops = to_uops_list([buf.shrink(((1, 3),)).bitcast(dtypes.uint64).index(0).store(1)], ren=Device[Device.DEFAULT].renderer)
    idx = next(u for u in uops if u.op is Ops.INDEX and u.src[0].op is Ops.BITCAST)
    self.assertEqual(idx.src[0].src[0].op, Ops.SHRINK)
@@ -71,7 +71,7 @@ class TestBitcastBufferView(unittest.TestCase):
  @Context(SPEC=2)
  def test_load(self):
    val = 0x1122334455667788
-    src, out = UOp.param(0, dtypes.uint32, 4), UOp.param(1, dtypes.uint64, 1)
+    src, out = UOp.param(0, dtypes.uint32, (4,)), UOp.param(1, dtypes.uint64, (1,))
    ibuf = Buffer(Device.DEFAULT, 4, dtypes.uint32, initial_value=np.array([0, 0x55667788, 0x11223344, 0], dtype=np.uint32).tobytes())
    obuf = Buffer(Device.DEFAULT, 1, dtypes.uint64).allocate()
    run_uops([out.index(0).store(src.shrink(((1, 3),)).bitcast(dtypes.uint64).index(0))], [ibuf, obuf])
@@ -80,7 +80,7 @@ class TestBitcastBufferView(unittest.TestCase):
  @Context(SPEC=2)
  def test_store(self):
    val = 0x1122334455667788
-    dst = UOp.param(0, dtypes.uint32, 6)
+    dst = UOp.param(0, dtypes.uint32, (6,))
    buf = Buffer(Device.DEFAULT, 6, dtypes.uint32, initial_value=bytes(24))
    view = dst.shrink(((1, 5),)).bitcast(dtypes.uint64)  # two stores through one view: it must inline, not get a declared vector-pointer
    run_uops([view.index(0).store(val ^ 0xff), view.index(1).store(val)], [buf])
@@ -249,8 +249,8 @@ class TestLocalAccess(unittest.TestCase):
@unittest.skipUnless(isinstance(Device[Device.DEFAULT].renderer, PTXRenderer), "This only tests assembly backends")
 class TestAssembly(unittest.TestCase):
  def test_bitshift_left(self):
-    g1 = UOp.param(0, dtypes.int32, 3)
-    out = UOp.param(1, dtypes.int32, 2)
+    g1 = UOp.param(0, dtypes.int32, shape=(3,))
+    out = UOp.param(1, dtypes.int32, shape=(2,))
    c1 = UOp.const(2)
    c2 = UOp.const(3)
    l1 = g1.index(c1)
@@ -278,7 +278,7 @@ class TestAssembly(unittest.TestCase):
    self.assertGreaterEqual(len([x.op for x in uops if x.op is Ops.MULACC]), 4)

  def test_mulacc_shl(self):
-    g1 = UOp.param(0, dtypes.int32, 2)
+    g1 = UOp.param(0, dtypes.int32, shape=(2,))
    c1 = UOp.const(0)
    c2 = UOp.const(1)
    expr = g1.index(c1) * UOp.const(4096) + g1.index(c2)
@@ -287,7 +287,7 @@ class TestAssembly(unittest.TestCase):
    self.assertIn(Ops.MULACC, [x.op for x in uops])

  def test_use_cmpeq(self):
-    g = UOp.param(0, dtypes.uint32, 8)
+    g = UOp.param(0, dtypes.uint32, shape=(8,))
    c = UOp.const(7)
    comp = g.index(c).ne(c).ne(True)
    uops = to_uops_list([comp], ren=Device[Device.DEFAULT].renderer)
@@ -43,7 +43,7 @@ def nested_loop_kernel(C:UOp) -> UOp:
  return C[0].store(i[0].load()).sink(arg=KernelInfo(name="nested_loop", opts_to_apply=()))

 def wait_ext_kernel() -> UOp:
-  sig = UOp.param(0, dtypes.int, 1, volatile=True)
+  sig = UOp.param(0, dtypes.int, (1,), volatile=True)
  l = UOp.loop(0)
  v = sig.after(l)[0].load()
  e = v.end(l, v < 1)
@@ -28,7 +28,7 @@ def vision_conv_143():
  c48 = (c24&c32).where(c34.index(c45), UOp.const(0.0, dtypes.float))
  c49 = UOp.param(2, dtypes.half, shape=(64, 49, 4))
  c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
-  c63 = UOp.param(3, dtypes.float, 128)
+  c63 = UOp.param(3, dtypes.float, (128,))
  c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
  c67 = c0.index((c2*128+c5+c8*4096)).store(c65).end(c8, c2, c5)

@@ -54,7 +54,7 @@ def vision_conv_153():
  c48 = (c24&c32).where(c34.index(c45), UOp.const(0.0, dtypes.float))
  c49 = UOp.param(2, dtypes.half, shape=(128, 49, 4))
  c61 = c48*c49.index((c26*4+c5%2+c16*28+c38*196))
-  c63 = UOp.param(3, dtypes.float, 256)
+  c63 = UOp.param(3, dtypes.float, (256,))
  c65 = c61.reduce(c16, c26, arg=Ops.ADD)+c63.index(c5)
  c67 = c0.index((c2*256+c5+c8*4096)).store(c65).end(c8, c2, c5)

@@ -73,7 +73,7 @@ def dm_conv_172():
  c18 = UOp.range(8, 2, AxisType.REDUCE)
  c23 = UOp.param(2, dtypes.half, shape=(240, 128, 4))
  c35 = c5.index((c7*4+c10+c13*128+c18*1536))*c23.index((c10*4+c2%4+c7*16+c2//4*512))
-  c37 = UOp.param(3, dtypes.float, 960)
+  c37 = UOp.param(3, dtypes.float, (960,))
  c39 = c35.reduce(c7, c10, arg=Ops.ADD)+c37.index(c2)
  c50 = (1.0+((c39+0.044708251953125*(c39*(c39*c39)))*-2.3021129851685216).exp2()).reciprocal()*c39
  c53 = c50.reduce(c18, c13, arg=Ops.ADD)*0.010416666666666666
@@ -36,7 +36,7 @@ class TestGPUCrash(unittest.TestCase):

  def _run_insts(self, insts: list[Inst]):
    buf = UOp.new_buffer("AMD", 64, dtypes.uint8)
-    sink = UOp.sink(UOp.param(0, dtypes.uint8, 64, device="AMD"), UOp.special(1, "lidx0"), arg=KernelInfo("test"))
+    sink = UOp.sink(UOp.param(0, dtypes.uint8, (64,), device="AMD"), UOp.special(1, "lidx0"), arg=KernelInfo("test"))
    prg = UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(UOp(Ops.INS, arg=i) for i in insts))))
    run_linear(UOp(Ops.LINEAR, src=(prg.call(buf),)), wait=True)

@@ -122,7 +122,7 @@ def eval_uop(uop:UOp, inputs:list[tuple[DType, list[Any]]]|None=None, vals:tuple
  for buf_dt, data in inputs or []:
    bufs.append(buf:=allocator.alloc(len(data) * buf_dt.itemsize))
    allocator._copyin(buf, memoryview(struct.pack(str(len(data)) + (buf_dt.fmt or ""), *data)))
-  g = UOp.param(0, uop.dtype, 1)
+  g = UOp.param(0, uop.dtype, (1,))
  prg = to_program(UOp.store(g.index(UOp.const(0)), uop).sink(arg=KernelInfo()), PythonRenderer(Target("PYTHON")))
  prog = dev.runtime(prg.to_elf())
  prog(out_buf:=allocator.alloc(uop.dtype.itemsize), *bufs, vals=vals)
@@ -321,10 +321,10 @@ def _int_clamp(op_name: str, srcs: dict) -> UOp | None:
 class _Ctx:
  """Context for instruction compilation - holds buffers and helpers."""
  __slots__ = ('inst_size', 'dyn_fields', '_axis_id', 'wave_size', 'vgpr', 'accvgpr')
-  sgpr = UOp.param(0, dtypes.uint32, SGPR_COUNT)
-  vmem = UOp.param(2, dtypes.uint32, 1 << 46)
-  lds = UOp.param(3, dtypes.uint32, 16384)
-  scratch = UOp.param(4, dtypes.uint8, 1 << 30)
+  sgpr = UOp.param(0, dtypes.uint32, (SGPR_COUNT,))
+  vmem = UOp.param(2, dtypes.uint32, (1 << 46,))
+  lds = UOp.param(3, dtypes.uint32, (16384,))
+  scratch = UOp.param(4, dtypes.uint8, (1 << 30,))
  # Cache PARAM UOps by wave_size so all _Ctx instances with same wave_size share identical UOp references
  _vgpr_cache: dict[int, UOp] = {}
  _accvgpr_cache: dict[int, UOp] = {}
@@ -332,10 +332,10 @@ class _Ctx:
  def __init__(self, inst_size: int, wave_size: int = 32):
    self.inst_size, self._axis_id, self.wave_size = inst_size, 0, wave_size
    self.dyn_fields: list[tuple[int, int]] = []  # (lo, hi) of fields read dynamically
-    if wave_size not in _Ctx._vgpr_cache: _Ctx._vgpr_cache[wave_size] = UOp.param(1, dtypes.uint32, 256 * wave_size)
+    if wave_size not in _Ctx._vgpr_cache: _Ctx._vgpr_cache[wave_size] = UOp.param(1, dtypes.uint32, (256 * wave_size,))
    self.vgpr = _Ctx._vgpr_cache[wave_size]
    if wave_size == 64:
-      if wave_size not in _Ctx._accvgpr_cache: _Ctx._accvgpr_cache[wave_size] = UOp.param(5, dtypes.uint32, 256 * wave_size)
+      if wave_size not in _Ctx._accvgpr_cache: _Ctx._accvgpr_cache[wave_size] = UOp.param(5, dtypes.uint32, (256 * wave_size,))
      self.accvgpr = _Ctx._accvgpr_cache[wave_size]
    else:
      self.accvgpr = self.vgpr
@@ -107,7 +107,7 @@ class TestGroupedDims(unittest.TestCase):

  def test_global_prod_max(self):
    g, l = UOp.range(256, 0, AxisType.GLOBAL), UOp.range(256, 1, AxisType.LOCAL)
-    sink = UOp.param(0, dtypes.float, 512).index(g + l).store(UOp.const(1.0)).end(g, l).sink(arg=KernelInfo())
+    sink = UOp.param(0, dtypes.float, (512,)).index(g + l).store(UOp.const(1.0)).end(g, l).sink(arg=KernelInfo())
    class R(Renderer): global_max, local_max, global_prod_max = (256, 256, 256), (128, 128, 128), (128, 128, 128)
    specials = [u for u in add_gpudims(R(Target()), sink).toposort() if u.op is Ops.SPECIAL]
    self.assertGreater(len([s for s in specials if "lidx" in s.arg]), 1)
@@ -7,14 +7,14 @@ from tinygrad.codegen import to_program

 class TestLinearizerFailures(unittest.TestCase):
  def test_fail_1(self):
-    c0 = UOp.param(0, dtypes.float, 64)
+    c0 = UOp.param(0, dtypes.float, (64,))
    c1 = UOp.range(UOp.const(2), 1, AxisType.WEAK)
    c2 = UOp.range(UOp.const(32), 2, AxisType.WEAK)
    c3 = ((c1*UOp.const(32))+c2)
-    c4 = UOp.param(1, dtypes.float, 163840)
+    c4 = UOp.param(1, dtypes.float, (163840,))
    c5 = UOp.range(UOp.const(2560), 0, AxisType.REDUCE)
    c6 = c4.index(((((((c5//UOp.const(8))%UOp.const(8))*UOp.const(8))+(c5%UOp.const(8)))+(((c2*UOp.const(40))+(c5//UOp.const(64)))*UOp.const(64)))+(c1*UOp.const(81920))))
-    c7 = UOp.param(2, dtypes.float, 64)
+    c7 = UOp.param(2, dtypes.float, (64,))
    c8 = c7.index(c3)
    c9 = ((((c6+(c8*UOp.const(-1.0)))*(c6+(c8*UOp.const(-1.0)))).reduce(c5, arg=Ops.ADD)*UOp.const(0.000390625))+UOp.const(1e-05)).sqrt().reciprocal()
    c10 = c0.index(c3).store(c9).end(c1, c2)
@@ -15,7 +15,7 @@ def simplify_valid_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move
 def simplify_image_idx(sink: UOp) -> UOp: return graph_rewrite(sink, sym+pm_move_where_on_load+indexing_simplify, name="simplify_image_idx")

 def get_gated_load_uop(valid:UOp, idx:UOp):
-  return UOp.param(0, dtypes.float, 1024).index(idx.valid(valid)).load()
+  return UOp.param(0, dtypes.float, (1024,)).index(idx.valid(valid)).load()

 def get_load_image_uop(image_shape:tuple[int, ...], valid:UOp, idx:tuple[UOp, UOp]):
  return UOp.param(0, dtypes.float, image_shape).index(idx[1].valid(valid), idx[0].valid(valid)).load()
@@ -505,7 +505,7 @@ class TestDropTrueGate(unittest.TestCase):
    from tinygrad.codegen.late.coalesce import indexing_simplify
    from tinygrad.uop.ops import graph_rewrite
    from tinygrad.uop.symbolic import sym
-    buf = UOp.param(0, dtypes.int, 1)
+    buf = UOp.param(0, dtypes.int, (1,))
    idx = UOp.const(0)
    true_gate = UOp.const(True)
    index_with_gate = buf.index(idx.valid(true_gate))
@@ -551,7 +551,7 @@ class TestRangeShrink(unittest.TestCase):
    # one load guards r < 4, but another load uses r without a gate -> no shrink
    r = Range(0, 204)
    load1 = get_gated_load_uop(r < UOp.const(4), r)
-    load2 = UOp.param(1, dtypes.float, 204).index(r).load()
+    load2 = UOp.param(1, dtypes.float, (204,)).index(r).load()
    ranges = self.get_ranges(UOp.sink(load1, load2))
    self.assert_range_end(ranges, 204)

@@ -575,7 +575,7 @@ class TestRangeShrink(unittest.TestCase):
    from tinygrad.dtype import Invalid
    r = Range(0, 204)
    x = (r < 4).where(UOp.const(1.0), Invalid)
-    ranges = self.get_ranges(UOp.param(0, dtypes.float, 204).index(r).store((r < 4).where(x, Invalid)).sink())
+    ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r < 4).where(x, Invalid)).sink())
    self.assert_range_end(ranges, 4)

  def test_range_shrink_store_where_invalid_flipped(self):
@@ -583,7 +583,7 @@ class TestRangeShrink(unittest.TestCase):
    from tinygrad.dtype import Invalid
    r = Range(0, 204)
    x = (r < 4).where(UOp.const(1.0), Invalid)
-    ranges = self.get_ranges(UOp.param(0, dtypes.float, 204).index(r).store((r >= 4).where(Invalid, x)).sink())
+    ranges = self.get_ranges(UOp.param(0, dtypes.float, (204,)).index(r).store((r >= 4).where(Invalid, x)).sink())
    self.assert_range_end(ranges, 4)

 if __name__ == '__main__':
@@ -9,7 +9,7 @@ class TestTranscendentalFunctions(unittest.TestCase):
  def test_payne_hanek_reduction(self):
    # TODO: Test constant input when constant folding is fixed (or maybe test both variants)
    # Load input value from a buffer to prevent constant folding
-    input_buf = UOp.param(1, dtypes.double, 1)
+    input_buf = UOp.param(1, dtypes.double, (1,))
    loaded_value = input_buf.index(UOp.const(0)).load()
    def eval_payne_hanek_reduction(v:float) -> tuple[float, int]:
      return tuple(eval_uop(u, [(dtypes.float64, [v])]) for u in payne_hanek_reduction(loaded_value))
@@ -214,8 +214,8 @@ class TestUOpGraph(unittest.TestCase):
      for i, const in enumerate(consts): self.assertIs(vec.index(i), const)

  def test_cast_alu_fold(self):
-    d0 = UOp.param(0, dtypes.bool, 1)
-    d1 = UOp.param(1, dtypes.int, 1)
+    d0 = UOp.param(0, dtypes.bool, (1,))
+    d1 = UOp.param(1, dtypes.int, (1,))
    idx = UOp.const(0)
    ld = d1.index(idx)
    alu = (ld<1).cast(dtypes.bool)
@@ -224,8 +224,8 @@ class TestUOpGraph(unittest.TestCase):
    self.assertEqual(len([x for x in uops if x.op is Ops.CAST and x.src[0].op is not Ops.CONST]), 0)

  def test_double_cast_fold(self):
-    d0 = UOp.param(0, dtypes.float, 1)
-    d1 = UOp.param(1, dtypes.int, 1)
+    d0 = UOp.param(0, dtypes.float, (1,))
+    d1 = UOp.param(1, dtypes.int, (1,))
    idx = UOp.const(0, dtypes.int)
    ld = d1.index(idx)
    alu = ld.cast(dtypes.float).cast(dtypes.float)
@@ -243,7 +243,7 @@ class TestUOpGraph(unittest.TestCase):

  def test_bitcast_to_same_dtype_fold(self):
    for dt in dtypes.ints + dtypes.floats + (dtypes.bool,):
-      d0 = UOp.param(0, dt, 1)
+      d0 = UOp.param(0, dt, (1,))
      v = d0.index(UOp.const(0))
      uops = to_uops_list([v.bitcast(dt)])
      self.assertEqual(len([x for x in uops if x.op is Ops.BITCAST and x.dtype is dt]), 0, f"dtype = {dt}")
@@ -255,10 +255,10 @@ class TestUOpGraph(unittest.TestCase):

  def test_where_on_gated_load_fold(self):
    ridx0 = UOp.range(100, 0)
-    d0 = UOp.param(0, dtypes.long, 100)
+    d0 = UOp.param(0, dtypes.long, (100,))
    ld = d0.index(ridx0.valid(ridx0<50))
    w = (ridx0<50).where(ld, 5)
-    out = UOp.param(1, dtypes.long, 100)
+    out = UOp.param(1, dtypes.long, (100,))
    uops = to_uops_list([out.index(ridx0).store(w)])
    expected = full_rewrite(UOp.const(5, dtypes.long).sink()).src[0]
    for u in uops:
@@ -267,7 +267,7 @@ class TestUOpGraph(unittest.TestCase):

  def test_where_on_gated_load_folds_swapped_branches(self):
    ridx0 = UOp.range(100, 0)
-    d0 = UOp.param(0, dtypes.long, 100)
+    d0 = UOp.param(0, dtypes.long, (100,))
    ld = d0.index(ridx0.valid((ridx0<50).logical_not()))
    w = (ridx0<50).where(5, ld)
    uops = to_uops_list([w])
@@ -278,11 +278,11 @@ class TestUOpGraph(unittest.TestCase):

  def test_where_on_gated_load_with_cast(self):
    ridx0 = UOp.range(100, 0)
-    d0 = UOp.param(0, dtypes.int, 100)
+    d0 = UOp.param(0, dtypes.int, (100,))
    gate_idx = ridx0.valid((ridx0<50))
    ld = d0.index(gate_idx).cast(dtypes.float)
    w = (ridx0<50).where(ld, 5.0)
-    out = UOp.param(1, dtypes.float, 100)
+    out = UOp.param(1, dtypes.float, (100,))
    uops = to_uops_list([out.index(ridx0).store(w)])
    expected = full_rewrite(UOp.const(5, dtypes.int).sink()).src[0]
    for u in uops:
@@ -291,27 +291,27 @@ class TestUOpGraph(unittest.TestCase):

  def test_where_on_casted_gated_load_extra_cond(self):
    ridx0 = UOp.range(100, 0)
-    d0 = UOp.param(0, dtypes.float, 100)
+    d0 = UOp.param(0, dtypes.float, (100,))
    ld = d0.index(ridx0.valid(ridx0<50))
    w = ((ridx0<50) & (ridx0>30)).where(ld, UOp.const(0.0)).cast(dtypes.half)
-    out = UOp.param(1, dtypes.half, 100)
+    out = UOp.param(1, dtypes.half, (100,))
    uops = to_uops_list([out.index(ridx0).store(w)])
    for u in uops:
      assert u.op is not Ops.WHERE

  def test_where_on_casted_gated_load_extra_cond_swapped(self):
    ridx0 = UOp.range(100, 0)
-    d0 = UOp.param(0, dtypes.float, 100)
+    d0 = UOp.param(0, dtypes.float, (100,))
    ld = d0.index(ridx0.valid(ridx0<50))
    w = ((ridx0<50) & (ridx0>30)).where(UOp.const(0.0), ld).cast(dtypes.half)
-    out = UOp.param(1, dtypes.half, 100)
+    out = UOp.param(1, dtypes.half, (100,))
    uops = to_uops_list([out.index(ridx0).store(w)])
    for u in uops:
      assert u.op is not Ops.WHERE

  def test_where_in_store_becomes_gate(self):
    ridx0 = UOp.range(100, 0)
-    d0 = UOp.param(0, dtypes.long, 100)
+    d0 = UOp.param(0, dtypes.long, (100,))
    idx = d0.index(ridx0)
    ld = idx.load()
    val = (ridx0<50).where(5, ld)
@@ -325,14 +325,14 @@ class TestUOpGraph(unittest.TestCase):
  def test_load_idx_becomes_int(self):
    # mnist indexing with split reduceop
    # Make sure we are not doign math on the loaded index, which would promote it to long
-    c0 = UOp.param(0, dtypes.uchar, 128000)
+    c0 = UOp.param(0, dtypes.uchar, (128000,))
    c1 = UOp.range(UOp.const(512), 1, AxisType.WEAK)
    c2 = UOp.range(UOp.const(250), 2, AxisType.WEAK)
-    c3 = UOp.param(1, dtypes.int, 512)
+    c3 = UOp.param(1, dtypes.int, (512,))
    c4 = c3.index(c1)
    c5 = UOp.range(UOp.const(240), 0, AxisType.REDUCE)
    c6 = ((c2*UOp.const(240))+c5)
-    c7 = UOp.param(2, dtypes.uchar, 60000)
+    c7 = UOp.param(2, dtypes.uchar, (60000,))
    c8 = c7.index(c6)
    c9 = ((c4<0).where((c4+60000), c4)!=c6.cast(dtypes.int)).where(0, c8.cast(dtypes.uint).cast(dtypes.uchar)).reduce(c5, arg=Ops.ADD)
    c10 = c0.index(((c1*UOp.const(250))+c2)).store(c9).end(c1, c2)
@@ -342,14 +342,14 @@ class TestUOpGraph(unittest.TestCase):

  def test_load_idx_no_math_on_loaded(self):
    # test the (x+y)<c pattern where x has loads - we shouldn't do math on loaded indices
-    c0 = UOp.param(0, dtypes.uchar, 128000)
+    c0 = UOp.param(0, dtypes.uchar, (128000,))
    c1 = UOp.range(UOp.const(512), 1, AxisType.WEAK)
    c2 = UOp.range(UOp.const(250), 2, AxisType.WEAK)
-    c3 = UOp.param(1, dtypes.int, 512)
+    c3 = UOp.param(1, dtypes.int, (512,))
    c4 = c3.index(c1)  # c4 is a load
    c5 = UOp.range(UOp.const(240), 0, AxisType.REDUCE)
    c6 = ((c2*UOp.const(240))+c5)
-    c7 = UOp.param(2, dtypes.uchar, 60000)
+    c7 = UOp.param(2, dtypes.uchar, (60000,))
    c8 = c7.index(c6)
    # (loaded + range) < const pattern - loaded value shouldn't be promoted to long
    loaded_idx = c4.cast(dtypes.weakint)
@@ -361,9 +361,9 @@ class TestUOpGraph(unittest.TestCase):
      self.assertNotEqual(u.dtype, dtypes.long)

  def test_fold_gated_load(self):
-    glbl0 = UOp.param(0, dtypes.int, 1)
-    glbl1 = UOp.param(1, dtypes.int, 1)
-    glbl2 = UOp.param(2, dtypes.int, 1)
+    glbl0 = UOp.param(0, dtypes.int, (1,))
+    glbl1 = UOp.param(1, dtypes.int, (1,))
+    glbl2 = UOp.param(2, dtypes.int, (1,))
    idx = UOp.const(0)
    ld0 = glbl1.index(UOp.invalid())
    ld1 = glbl2.index(idx.valid(UOp.const(True)))
@@ -372,7 +372,7 @@ class TestUOpGraph(unittest.TestCase):
    self.assertEqual(len([u for u in uops if u.op is Ops.LOAD]), 1)

  def test_fold_gated_load_local(self):
-    glbl0 = UOp.param(0, dtypes.int, 16)
+    glbl0 = UOp.param(0, dtypes.int, (16,))
    smem = UOp.placeholder((18,), dtypes.int, slot=0, addrspace=AddrSpace.LOCAL)
    lidx = UOp.special(16, "lidx0")
    st = smem.index(lidx).store(glbl0.index(lidx).load())
@@ -385,7 +385,7 @@ class TestUOpGraph(unittest.TestCase):
    self.assertEqual(len([u for u in uops if u.op is Ops.LOAD]), 2)

  def test_fold_gated_store(self):
-    glbl = UOp.param(0, dtypes.int, 1)
+    glbl = UOp.param(0, dtypes.int, (1,))
    idx0 = UOp.const(0)
    val = UOp.const(42)
    st0 = glbl.index(UOp.invalid()).store(val)
@@ -425,7 +425,7 @@ class TestReduceCollapse(unittest.TestCase):

  def test_reduce_shapeless_const_unroll(self):
    """a REDUCE over a shapeless CONST (e.g. x*0 folded late in codegen) must collapse before the expander"""
-    out = UOp.param(0, dtypes.float, 1)
+    out = UOp.param(0, dtypes.float, (1,))
    red = UOp.const(3.0).cast(dtypes.float).reduce(UOp.range(4, 0, AxisType.UNROLL), arg=(Ops.ADD, 0))
    ast = UOp.sink(out.index(UOp.const(0)).store(red)).replace(arg=KernelInfo())
    uops = full_rewrite_to_sink(ast, Device["CPU"].renderer, optimize=False).toposort()
@@ -441,8 +441,7 @@ class TestMovementOps(unittest.TestCase):
    self.assertEqual(result.op, Ops.INDEX)
    self.assertIs(result.src[0], src)
    self.assertEqual(result.shape, (4,))
-    # the only RESHAPE is src itself: the view of the flat param, the extra reshape was folded into the INDEX
-    self.assertEqual([u for u in result.toposort() if u.op is Ops.RESHAPE], [src])
+    self.assertNotIn(Ops.RESHAPE, [u.op for u in result.toposort()])

  def test_pm_mops_partial_reshape_index_suffix_mismatch_does_nothing(self):
    from tinygrad.schedule.prepare import pm_mops
@@ -1392,7 +1392,7 @@ class TestInvalidIndex(unittest.TestCase):

  def test_gated_load_keeps_index_valid(self):
    # the load executes even on gated-off iterations: gated_given_valid must not erase its mask (PADTO OOB shape)
-    buf = UOp.param(0, dtypes.bool, 17)
+    buf = UOp.param(0, dtypes.bool, (17,))
    ridx = Variable("ridx", 0, 31)
    cond = ridx < 17
    load = buf.index(ridx.valid(cond))
@@ -1404,7 +1404,7 @@ class TestStoreLoadFolding(unittest.TestCase):
  """Tests for store(index, load(index)) -> NOOP rule. This rule matches patterns that EMERGE during simplification."""
  def test_store_load_folding(self):
    # store(idx, load(idx)) -> NOOP, including emergent patterns like store(idx, load(idx) + 0)
-    buf = UOp.param(0, dtypes.int, 1)
+    buf = UOp.param(0, dtypes.int, (1,))
    index = buf.index(UOp.const(0))
    # Direct: store(idx, load(idx)) -> NOOP
    self.assertEqual(graph_rewrite(index.store(index.load()), sym).op, Ops.NOOP)
@@ -1417,7 +1417,7 @@ class TestStoreLoadFolding(unittest.TestCase):

 class TestMoveWhereOnLoad(unittest.TestCase):
  def test_bool_index_preserves_dtype(self):
-    buf = UOp.param(0, dtypes.bool, 8)
+    buf = UOp.param(0, dtypes.bool, (8,))
    a = Variable("a", 0, 7)
    r = UOp.range(8, 0)
    # cond has a range that the rewrite can move into the valid: gate (a<4) goes into load valid
@@ -1475,7 +1475,7 @@ class TestRangeSplitting(unittest.TestCase):
    from tinygrad.codegen.simplify import pm_split_ranges, pm_flatten_range
    r0 = UOp.range(uconst(8), 0)
    # create a simple expression using the range with mod: store range%2 to a buffer
-    buf = UOp.param(0, dtypes.int, 1)
+    buf = UOp.param(0, dtypes.int, (1,))
    val = (r0 % uconst(2)).cast(dtypes.int)
    store = UOp(Ops.STORE, src=(buf.index(uconst(0)), val))
    sink = UOp(Ops.SINK, src=(UOp(Ops.END, src=(store, r0)),))
@@ -82,7 +82,7 @@ class TestVminVmaxProperties(unittest.TestCase):
  def test_vmin_vmax_multiplication_0_inf(self):
    # vmin and vmax for multiplication with a variable
    x = UOp.const(0.0)
-    y = UOp.load(UOp.param(0, dtypes.float, 1), UOp.const(0))
+    y = UOp.load(UOp.param(0, dtypes.float, (1,)), UOp.const(0))
    uop = x * y
    # TODO: these should be 0, but definitely should not be nan
    self.assertEqual(uop.vmin, -math.inf)
@@ -332,7 +332,7 @@ class TestVminVmaxVConst(unittest.TestCase):

  def test_vmin_vmax_vector_with_gep(self):
    # vmin and vmax for a vector constant of bool values
-    d1 = UOp.param(1, dtypes.int, 1)
+    d1 = UOp.param(1, dtypes.int, (1,))
    idx = UOp.const(0)
    val = UOp(Ops.LOAD, src=(d1.index(idx),))
    uop = (val // 32)
@@ -58,7 +58,7 @@ class TestDTypeFromUOp(unittest.TestCase):
    for u in (UOp.param(0, dtypes.bool, ()).where(value, invalid), value+invalid, UOp.stack(value, invalid)): self.assertIs(u.src[-1], invalid)
    for u in (UOp(Ops.STACK, src=(value, invalid)), UOp(Ops.ADD, src=(value, invalid)),
              UOp.const(True).where(value, invalid), UOp(Ops.CMPLT, src=(invalid, value)), UOp(Ops.CMPLT, src=(value, invalid)),
-              UOp.param(0, dtypes.float32, 4).index(invalid)): type_verify(u, spec_shared)
+              UOp.param(0, dtypes.float32, (4,)).index(invalid)): type_verify(u, spec_shared)
    gate, value = UOp.param(0, dtypes.bool, ()), UOp.param(1, dtypes.float, ())
    self.assertIs((out:=graph_rewrite(gate.where(value, UOp.invalid()), pm_remove_invalid)).src[2], UOp.const(0, dtypes.float))
    type_verify(out.sink(), spec_program)
@@ -73,7 +73,7 @@ class TestLowerIndexDtype(unittest.TestCase):
  def test_gated_shrink_lowers_to_selected_width(self):
    # coalesce builds gated SHRINKs for masked vectorized loads; lowering must resolve them at the
    # width the offset bounds select (this one needs long)
-    buf = UOp.param(0, dtypes.float, 2**31+64)
+    buf = UOp.param(0, dtypes.float, (2**31+64,))
    i = UOp.variable("i", 0, 2**28)
    shrink = UOp(Ops.SHRINK, src=(buf, (i*24).valid(i < 2**28), UOp.const(4)))
    lowered = graph_rewrite(shrink.sink(), pm_lower_weak)
@@ -84,10 +84,11 @@ class TestLowerIndexDtype(unittest.TestCase):

  def test_reg_buffer_size_lowers(self):
    reg = UOp.placeholder((4,), dtypes.float, 0, addrspace=AddrSpace.REG)
-    self.assertEqual(reg.arg.size, 4)
+    self.assertEqual(reg.src[0].dtype, dtypes.weakint)
    lowered = graph_rewrite(reg.sink(), pm_lower_weak)
    self.assertTrue(all(u.op is Ops.CONST for u in lowered.backward_slice_with_self if u.dtype in dtypes.weaks),
                    "lowering must resolve every weak width, except a typed literal's value half")
+    self.assertEqual(next(u for u in lowered.backward_slice_with_self if u.op is Ops.BUFFER).src[0].dtype, dtypes.int)

 class TestSafeCast(unittest.TestCase):
  def test_cast_folds(self):
@@ -220,7 +221,7 @@ class TestExecALU(unittest.TestCase):

 class TestGatedStoreRewrite(unittest.TestCase):
  def test_tiny_gate_store(self):
-    gmem = UOp.param(0, dtypes.float, 8)
+    gmem = UOp.param(0, dtypes.float, (8,))
    gidx0 = UOp.special(4, 'gidx0')
    gate = gidx0<UOp.const(1)
    idx = UOp(Ops.INDEX, src=(gmem, (gidx0 * UOp.const(2)).valid(gate)))
@@ -236,8 +237,8 @@ class TestGatedStoreRewrite(unittest.TestCase):
    self.assertEqual(len(gated_uops[-1].src), 2)

  def test_gate_some_stores(self):
-    gmem0 = UOp.param(0, dtypes.float, 8)
-    gmem1 = UOp.param(1, dtypes.float, 8)
+    gmem0 = UOp.param(0, dtypes.float, (8,))
+    gmem1 = UOp.param(1, dtypes.float, (8,))
    gidx0 = UOp.special(4, 'gidx0')
    idx = gidx0 * UOp.const(2)
    idx0 = UOp(Ops.INDEX, src=(gmem0, idx.valid(gidx0<UOp.const(1))))
@@ -256,8 +257,8 @@ class TestGatedStoreRewrite(unittest.TestCase):
  # scaled down version of TestLinearizerDumb.test_unmerged_ifs
  @unittest.skip("we don't merge ifs anymore")
  def test_merge_ifs_alt(self):
-    gmem0 = UOp.param(0, dtypes.float, 8)
-    gmem1 = UOp.param(1, dtypes.float, 8)
+    gmem0 = UOp.param(0, dtypes.float, (8,))
+    gmem1 = UOp.param(1, dtypes.float, (8,))
    gidx0 = UOp.special(4, 'gidx0')
    idx = gidx0*UOp.const(2)
    gate = gidx0<UOp.const(1)
@@ -280,7 +281,7 @@ class TestGatedStoreRewrite(unittest.TestCase):
 class TestFastIdiv(unittest.TestCase):
  def test_division_power_of_two(self):
    for dt in (dtypes.int32, dtypes.uint32):
-      g = UOp.param(0, dt, 3)
+      g = UOp.param(0, dt, (3,))
      c = UOp.const(2)
      l = g.index(c)
      a = UOp(Ops.CDIV, src=(l, c))
@@ -293,7 +294,7 @@ class TestFastIdiv(unittest.TestCase):
  def test_floormod_power_of_two(self):
    # FLOORMOD by a power of two lowers to AND (correct floor mod for any sign in two's complement)
    for dt in (dtypes.int32, dtypes.uint32):
-      g = UOp.param(0, dt, 9)
+      g = UOp.param(0, dt, (9,))
      c = UOp.const(8)
      a = UOp(Ops.FLOORMOD, src=(g.index(c), c))
      uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
@@ -305,7 +306,7 @@ class TestFastIdiv(unittest.TestCase):
  def test_floordiv_power_of_two(self):
    # FLOORDIV by a power of two lowers to a shift, with no round toward zero correction (a shift is exactly floor division)
    for dt in (dtypes.int32, dtypes.uint32, dtypes.int64, dtypes.uint64):
-      g = UOp.param(0, dt, 3)
+      g = UOp.param(0, dt, (3,))
      c = UOp.const(2)
      a = UOp(Ops.FLOORDIV, src=(g.index(c), c))
      uops = to_uops_list([a], ren=Device[Device.DEFAULT].renderer)
@@ -318,7 +319,7 @@ class TestFastIdiv(unittest.TestCase):
  @Context(DISABLE_FAST_IDIV=0)
  @unittest.skipIf(Device.DEFAULT == "WEBGPU", "WEBGPU doesn't support long")
  def test_fast_idiv_and_mod(self):
-    g = UOp.param(0, dtypes.uint32, 4)
+    g = UOp.param(0, dtypes.uint32, (4,))
    c = UOp.const(3)
    l = g.index(c)
    a = UOp(Ops.CDIV, src=(l, c))
@@ -352,7 +353,7 @@ class TestFastIdiv(unittest.TestCase):
  @unittest.expectedFailure
  def test_fast_idiv_overflow(self):
    # This will be possible with a slightly different method for fast_idiv
-    g = UOp.param(0, dtypes.uint32, 8)
+    g = UOp.param(0, dtypes.uint32, (8,))
    c = UOp.const(7).cast(dtypes.uint)
    l = UOp(Ops.LOAD, src=(g.index(c),))
    a = UOp(Ops.CDIV, src=(l, c))
@@ -363,7 +364,7 @@ class TestFastIdiv(unittest.TestCase):
    self.assertNotIn(Ops.CDIV, ops)

  def test_disable_fast_idiv(self):
-    g = UOp.param(0, dtypes.uint32, 4)
+    g = UOp.param(0, dtypes.uint32, (4,))
    c = UOp.const(3)
    l = g.index(c)
    a = UOp(Ops.CDIV, src=(l, c))
@@ -406,8 +407,8 @@ class TestUOpMethod(unittest.TestCase):
      self.assertIs((x != x).simplify(), x.const_like(False, dtypes.bool))

  def test_replace(self):
-    x = UOp.param(0, dtypes.int, 1)
-    self.assertEqual(x.replace(arg=UOp.param(1, dtypes.int, 1).arg).arg.slot, 1)
+    x = UOp.param(0, dtypes.int, (1,))
+    self.assertEqual(x.replace(arg=UOp.param(1, dtypes.int, (1,)).arg).arg.slot, 1)
    with self.assertRaises(AssertionError): x.replace(field="a")

  def test_const_zero_neg_zero_different(self):
@@ -462,7 +463,7 @@ class TestUOpRender(unittest.TestCase):
  def test_render_ssimplified_marg_outside_toposort(self):
    r = UOp.range(UOp.const(16, dtypes.int), 2, AxisType.WEAK, dtype=dtypes.int)
    offset = (r * 2) + (r * 2)
-    shrink = UOp(Ops.SHRINK, src=(UOp.param(0, dtypes.uint, 32), offset, UOp.const(2, dtypes.int)))
+    shrink = UOp(Ops.SHRINK, src=(UOp.param(0, dtypes.uint, (32,)), offset, UOp.const(2, dtypes.int)))
    self.assertIsNot(shrink.src[1], shrink.marg[0][0])
    self.assertEqual(shrink.render(simplify=False), "p0.shrink((((r2*4), 2),))")
    self.assertEqual(UOp.range(1, 0, src=(shrink,), dtype=dtypes.int).render(simplify=False), "r0")
@@ -146,7 +146,7 @@ class TestUOpsStats(unittest.TestCase):

  #MULACC should have the same stats as MUL + ADD
  def test_mulacc(self):
-    globl = UOp.param(0, dtypes.int, 3)
+    globl = UOp.param(0, dtypes.int, (3,))
    o1 = UOp.const(1, dtypes.int)
    o2 = UOp.const(2, dtypes.int)
    u1 = globl.index(o1)
@@ -156,7 +156,7 @@ class TestUOpsStats(unittest.TestCase):
    u5 = UOp(Ops.ADD, src=(u4,u3))
    uops = tuple(u5.toposort())

-    globl = UOp.param(0, dtypes.int, 3)
+    globl = UOp.param(0, dtypes.int, (3,))
    o1 = UOp.const(1, dtypes.int)
    o2 = UOp.const(2, dtypes.int)
    u1 = globl.index(o1)
@@ -13,7 +13,7 @@ class TestValidateOOB(unittest.TestCase):
  # basic index patterns
  def test_const_index(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, 16)
+      buf = UOp.param(0, dtypes.int, (16,))
      to_uops_list([buf.index(UOp.const(0)).load()])  # valid
      to_uops_list([buf.index(UOp.const(15)).load()])  # valid (last element)
      with self.assertRaises(RuntimeError):
@@ -23,7 +23,7 @@ class TestValidateOOB(unittest.TestCase):

  def test_variable_index(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, 16)
+      buf = UOp.param(0, dtypes.int, (16,))
      to_uops_list([buf.index(Variable("i", 0, 15)).load()])  # valid
      with self.assertRaises(RuntimeError):
        to_uops_list([buf.index(Variable("i", 0, 20)).load()])  # oob
@@ -32,7 +32,7 @@ class TestValidateOOB(unittest.TestCase):

  def test_range_with_mask(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, 16)
+      buf = UOp.param(0, dtypes.int, (16,))
      r = UOp.range(42, 0, AxisType.GLOBAL)
      to_uops_list([buf.index(r.valid(r < 16)).load()])  # valid
      with self.assertRaises(RuntimeError):
@@ -40,7 +40,7 @@ class TestValidateOOB(unittest.TestCase):

  def test_variable_with_mask(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, 16)
+      buf = UOp.param(0, dtypes.int, (16,))
      v = Variable("v", -5, 80)
      to_uops_list([buf.index(v.valid((v >= 0) & (v < 16))).load()])  # valid
      with self.assertRaises(RuntimeError):
@@ -48,7 +48,7 @@ class TestValidateOOB(unittest.TestCase):

  def test_gated_store(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, 16)
+      buf = UOp.param(0, dtypes.int, (16,))
      v = Variable("v", 0, 20)
      to_uops_list([buf.index(v.valid(v < 16)).store(0)])  # valid
      with self.assertRaises(RuntimeError):
@@ -57,14 +57,14 @@ class TestValidateOOB(unittest.TestCase):
  # ALU ops in index
  def test_floordiv(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, 16)
+      buf = UOp.param(0, dtypes.int, (16,))
      to_uops_list([buf.index(UOp.range(32, 0, AxisType.GLOBAL) // 2).load()])  # 0..15 valid
      with self.assertRaises(RuntimeError):
        to_uops_list([buf.index(UOp.range(34, 0, AxisType.GLOBAL) // 2).load()])  # 0..16 oob

  def test_mod(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, 16)
+      buf = UOp.param(0, dtypes.int, (16,))
      r = UOp.range(100, 0, AxisType.GLOBAL)
      to_uops_list([buf.index(r % 16).load()])  # 0..15 valid
      with self.assertRaises(RuntimeError):
@@ -72,14 +72,14 @@ class TestValidateOOB(unittest.TestCase):

  def test_shr(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, 16)
+      buf = UOp.param(0, dtypes.int, (16,))
      to_uops_list([buf.index(UOp.range(64, 0, AxisType.GLOBAL) >> 2).load()])  # 0..15 valid
      with self.assertRaises(RuntimeError):
        to_uops_list([buf.index(UOp.range(128, 0, AxisType.GLOBAL) >> 2).load()])  # 0..31 oob

  def test_shl(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, 64)
+      buf = UOp.param(0, dtypes.int, (64,))
      r = UOp.range(8, 0, AxisType.GLOBAL)
      to_uops_list([buf.index(r << 2).load()])  # 0..28 valid
      with self.assertRaises(RuntimeError):
@@ -87,7 +87,7 @@ class TestValidateOOB(unittest.TestCase):

  def test_and(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, 16)
+      buf = UOp.param(0, dtypes.int, (16,))
      r = UOp.range(100, 0, AxisType.GLOBAL)
      to_uops_list([buf.index(r & 15).load()])  # 0..15 valid
      with self.assertRaises(RuntimeError):
@@ -102,14 +102,14 @@ class TestValidateOOB(unittest.TestCase):

  def test_max(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, 16)
+      buf = UOp.param(0, dtypes.int, (16,))
      to_uops_list([buf.index(Variable("v", -10, 15).maximum(0)).load()])  # 0..15 valid
      with self.assertRaises(RuntimeError):
        to_uops_list([buf.index(Variable("v2", -10, 20).maximum(0)).load()])  # 0..20 oob

  def test_xor_in_mask(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, 16)
+      buf = UOp.param(0, dtypes.int, (16,))
      r = UOp.range(32, 0, AxisType.GLOBAL)
      to_uops_list([buf.index(r.valid((r < 8) ^ ((r >= 8) & (r < 16)))).load()])  # 0..15 valid
      with self.assertRaises(RuntimeError):
@@ -118,22 +118,22 @@ class TestValidateOOB(unittest.TestCase):
  # cast patterns
  def test_float_cast_in_index(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, 16)
+      buf = UOp.param(0, dtypes.int, (16,))
      r = UOp.range(20, 0)
      i = (r.cast(dtypes.float) * 0.68).trunc().cast(dtypes.int)
      to_uops_list([buf.index(i.valid((i >= 0) & (i < 16))).load()])

  def test_bool_cast_in_mask(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf = UOp.param(0, dtypes.int, 1)
+      buf = UOp.param(0, dtypes.int, (1,))
      r = UOp.range(20, 0)
      to_uops_list([buf.index(r.valid(r.cast(dtypes.bool).logical_not())).load()])  # only r=0 valid

  # load result as index/mask
  def test_load_as_index(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf0 = UOp.param(0, dtypes.int, 16)
-      buf1 = UOp.param(1, dtypes.int, 64)
+      buf0 = UOp.param(0, dtypes.int, (16,))
+      buf1 = UOp.param(1, dtypes.int, (64,))
      r = UOp.range(42, 0, AxisType.GLOBAL)
      ld0 = buf0.index(r.valid(r < 8)).load().cast(dtypes.weakint)
      to_uops_list([buf1.index((ld0 * 2).valid((ld0 >= 0) & (ld0 < 32))).load()])  # valid
@@ -142,16 +142,16 @@ class TestValidateOOB(unittest.TestCase):

  def test_load_from_shrink_as_index(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf0 = UOp.param(0, dtypes.int, 16)
-      buf1 = UOp.param(1, dtypes.int, 64)
+      buf0 = UOp.param(0, dtypes.int, (16,))
+      buf1 = UOp.param(1, dtypes.int, (64,))
      shrink = UOp(Ops.SHRINK, src=(buf0, UOp.const(0, dtypes.int), UOp.const(4)))
      ld0 = shrink.load().index(0)
      to_uops_list([buf1.index(ld0.valid((ld0 >= 0) & (ld0 < 64))).load()])

  def test_load_bool_as_mask(self):
    with Context(CHECK_OOB=1, SPEC=2):
-      buf_bool = UOp.param(0, dtypes.bool, 16)
-      buf_int = UOp.param(1, dtypes.int, 8)
+      buf_bool = UOp.param(0, dtypes.bool, (16,))
+      buf_int = UOp.param(1, dtypes.int, (8,))
      gidx = UOp(Ops.SPECIAL, src=(UOp.const(16),), arg="gidx0")
      ld_bool = buf_bool.index(gidx).load()
      with self.assertRaises(RuntimeError):
@@ -162,7 +162,7 @@ class TestValidateOOB(unittest.TestCase):
  def test_in_bounds_access_gated_local(self):
    with Context(CHECK_OOB=1):
      # Define buffers
-      gbuf = UOp.param(0, dtypes.uint, 400)
+      gbuf = UOp.param(0, dtypes.uint, (400,))
      sbuf = UOp.placeholder((8,), dtypes.uint, slot=0, addrspace=AddrSpace.LOCAL)

      # Define indices, valids and barrier
@@ -186,8 +186,8 @@ class TestValidateOOB(unittest.TestCase):
  @unittest.skip("Bool load is not supported yet")
  def test_load_mask(self):
    with Context(CHECK_OOB=1):
-      glbl0 = UOp.param(0, dtypes.int, 16)
-      mask = UOp.param(0, dtypes.bool, 16)
+      glbl0 = UOp.param(0, dtypes.int, (16,))
+      mask = UOp.param(0, dtypes.bool, (16,))
      ridx = UOp.range(20, 0)
      ld0 = UOp(Ops.LOAD, src=(glbl0.index(UOp.const(ridx<16&mask, ridx))))
      to_uops_list([ld0])
@@ -244,7 +244,7 @@ class TestViz(unittest.TestCase):
    self.assertEqual(len(lst), 1)
    graphs = [x["graph"] for x in viz.get_details(0, 0)]
    # const is always in the graph, client side hides exclude=True nodes by default
-    self.assertEqual(list(graphs[0]), [id(a), id(z), id(alu), id(y), id(sink)])
+    self.assertEqual(list(graphs[0]), [id(a.src[0]), id(a), id(z), id(alu), id(y), id(sink)])
    self.assertTrue(graphs[0][id(z)]["exclude"])
    self.assertTrue(graphs[0][id(y)]["exclude"])
    self.assertFalse(graphs[0][id(alu)]["exclude"])
@@ -73,10 +73,10 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
    return Tensor.linspace(start, stop, int(np.prod(shape)), dtype=dtypes.float32).reshape(*shape)

  def _make_config(self, **kwargs):
-    return TransformerConfig(**({"num_blocks":1, "dim":8, "hidden_dim":16, "n_heads":1, "n_kv_heads":1,
-                                 "norm_eps":1e-5, "vocab_size":32, "head_dim":8, "rope_theta":10000.0,
-                                 "rope_dim":8, "v_head_dim":8, "max_context":4, "ssm_layers":(True,),
-                                 "ssm":SSMConfig(conv_kernel=2, state_size=4, group_count=1, time_step_rank=1, inner_size=4)} | kwargs))
+    return TransformerConfig(**({"num_blocks":1, "dim":32, "hidden_dim":64, "n_heads":1, "n_kv_heads":1,
+                                 "norm_eps":1e-5, "vocab_size":32, "head_dim":32, "rope_theta":10000.0,
+                                 "rope_dim":32, "v_head_dim":32, "max_context":4, "ssm_layers":(True,),
+                                 "ssm":SSMConfig(conv_kernel=2, state_size=32, group_count=1, time_step_rank=1, inner_size=32)} | kwargs))

  def _make_block(self, config:TransformerConfig) -> GatedDeltaNetBlock:
    block = GatedDeltaNetBlock(config, config.ssm)
@@ -229,7 +229,7 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
    np.testing.assert_allclose(block.recurrent_state.numpy(), initial_state.numpy() * alpha[..., None], rtol=1e-5, atol=1e-5)

  def test_kda_prefill_matches_decode(self):
-    config = self._make_config(ssm=SSMConfig(conv_kernel=2, state_size=4, group_count=1, time_step_rank=1, inner_size=4, kda=True))
+    config = self._make_config(ssm=SSMConfig(conv_kernel=2, state_size=32, group_count=1, time_step_rank=1, inner_size=32, kda=True))
    block = GatedDeltaNetBlock(config, config.ssm)
    for p in nn.state.get_parameters(block):
      p.replace(self._tensor_linspace(-0.05, 0.05, p.shape) if len(p.shape) > 1 else self._tensor_linspace(0.05, 0.1, p.shape))
@@ -245,7 +245,7 @@ class TestGatedDeltaNetBlock(unittest.TestCase):

  def test_varied_chunk_sizes_match_decode(self):
    for kda in (False, True):
-      ssm = SSMConfig(conv_kernel=2, state_size=4, group_count=1, time_step_rank=1, inner_size=4, kda=kda)
+      ssm = SSMConfig(conv_kernel=2, state_size=32, group_count=1, time_step_rank=1, inner_size=32, kda=kda)
      config = self._make_config(ssm=ssm)
      if kda:
        block = GatedDeltaNetBlock(config, config.ssm)
@@ -267,8 +267,7 @@ class TestGatedDeltaNetBlock(unittest.TestCase):
        np.testing.assert_allclose(chunked_recurrent, decode_recurrent, rtol=1e-3, atol=1e-3, err_msg=f"{kda=} {chunking=}")

  def test_start_zero_resets_realized_state(self):
-    config = self._make_config(max_context=3)
-    x = self._tensor_linspace(-1, 1, (1, 3, config.dim))
+    config, x = self._make_config(max_context=3), self._tensor_linspace(-1, 1, (1, 3, 32))
    block = self._make_block(config)
    self._run_attention(block, x, 0)
    restarted = self._run_attention(block, x[:, :2], 0)
@@ -52,11 +52,6 @@ class TestCall(unittest.TestCase):
    np.testing.assert_allclose(a.grad.numpy(), gt_a_grad, rtol=1e-5)
    np.testing.assert_allclose(b.grad.numpy(), gt_b_grad, rtol=1e-5)

-  def test_call_scalar_param_shape_mismatch(self):
-    scalar_fxn = UOp.param(0, dtypes.float, ()) * 2
-    with self.assertRaisesRegex(TypeError, "shape mismatch: expected scalar"):
-      Tensor.call(Tensor.ones(2), fxn=scalar_fxn).realize()
-
  def test_call_gemm(self):
    M, K, N = 4, 8, 4
    a = Tensor.randn(M, K)
@@ -116,7 +116,7 @@ class TestWeakPromotion(unittest.TestCase):

  def test_store_weak_value_uses_destination_dtype(self):
    with Context(DEFAULT_FLOAT=dtypes.float16):
-      dst = UOp.param(0, dtypes.bfloat16, 1).index(UOp.const(0).cast(dtypes.int32))
+      dst = UOp.param(0, dtypes.bfloat16, (1,)).index(UOp.const(0).cast(dtypes.int32))
      gate = UOp.const(True)
      out = graph_rewrite(dst.store(UOp.const(5.0), gate), pm_commit_weak)
    # a bare weak CONST commits directly: the pass runs without symbolic, so a CAST here would survive it
@@ -132,7 +132,7 @@ class TestWeakPromotion(unittest.TestCase):

  def test_derivable_const_rounds_at_the_derived_width(self):
    # re-rounds a derivable const in place (still bare) so value-keyed folds (x*1 -> x, x*-1 -> NEG) still fire
-    x = UOp.param(0, dtypes.float32, 1).index(UOp.const(0).cast(dtypes.int32)).load()
+    x = UOp.param(0, dtypes.float32, (1,)).index(UOp.const(0).cast(dtypes.int32)).load()
    mul = graph_rewrite(x * UOp.const(-0.9999999893980771), symbolic_simple+pm_commit_weak)
    self.assertIs(mul.src[1], UOp.const(-1.0))
    self.assertIs(graph_rewrite(x * UOp.const(1.0000000106), symbolic_simple+pm_commit_weak), x)
@@ -391,15 +391,7 @@ def full_rewrite_to_sink(ast:UOp, ren:Renderer, optimize:bool=True) -> UOp:
  sink = graph_rewrite(sink, pm_number_params, ctx=[num_params], name="number params with -1", walk=True)

  if VIZ: graph_rewrite(sink, PatternMatcher([]), name="View Output AST")
-  if SPEC:
-    import os
-    if os.environ.get("DBGTV"):
-      try: type_verify(sink, spec_program)
-      except RuntimeError:
-        from tinygrad.uop.render import print_uops
-        print_uops(list(sink.toposort()))
-        raise
-    else: type_verify(sink, spec_program)
+  if SPEC: type_verify(sink, spec_program)

  # return the rewritten sink
  return sink
@@ -140,8 +140,8 @@ def f2f_store(st, idx, val, fr:DType, to:DType):
 pm_long_decomp: PatternMatcher = PatternMatcher([
  # the decomp's own bottom-up rewrite can mint bare consts mid-flight: word splitting commits them at the long sibling's dtype
  (UPat(GroupOp.All, name='x'), lambda x: commit_weak_consts(x, next((s.dtype for s in x.src if s.dtype in l2i_dt), None))),
-  (UPat(GroupOp.Defines, tuple(l2i_dt.keys()), name="x"), lambda x:
-   UOp(x.op, arg=replace(x.arg, dtype=l2i_dt[x.dtype], size=None if x.arg.size is None else x.arg.size*2), tag=x.tag)),
+  (UPat(GroupOp.Defines, tuple(l2i_dt.keys()), src=(UPat.var("sz"),), name="x"), lambda x,sz:
+   UOp(x.op, src=(sz*2,), arg=replace(x.arg, dtype=l2i_dt[x.dtype]), tag=x.tag)),
  (UPat(Ops.INDEX, tuple(l2i_dt.keys()), name='x'), lambda x:
   reindex(x, x.tag[0]).replace(tag=None) if x.tag is not None else None),
  (UPat(Ops.STORE, src=(UPat.var('idx', tuple(l2i_dt.keys())), UPat.var('val')), name='st'), lambda st,idx,val:
@@ -73,7 +73,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 len((idx := r.src[0]).src) and idx.src[0].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"
@@ -1,8 +1,7 @@
 import itertools, functools
 from collections import defaultdict
-from dataclasses import replace
 from tinygrad.dtype import dtypes, AddrSpace, Invalid, DType
-from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat, GroupOp, graph_rewrite
+from tinygrad.uop.ops import UOp, Ops, PatternMatcher, UPat, GroupOp, shape_to_shape_arg, graph_rewrite
 from tinygrad.uop.symbolic import uop_given_valid, parse_valid, invalid_gate, sym
 from tinygrad.helpers import getenv, IMAGE, OSX, ceildiv, is_image_shape
 from tinygrad.renderer import Renderer
@@ -86,8 +85,7 @@ def transform_to_image(ctx, buf:UOp, x:UOp) -> UOp|None:
  if len(cands) == 0: return None
  # and tiebreak with indexing complexity (ie. number of nodes)
  h, w, cidx = cands[0] if len(cands) == 1 else min(cands, key=lambda cand: len(cand[2].index(1).simplify().backward_slice))
-  # the image dims are stored in the param's arg, the size stays the flat buffer len
-  buf = buf.replace(arg=replace(buf.arg, image=(h, w)))
+  buf = buf.replace(src=(shape_to_shape_arg((h, w, 4)),))
  shapes[buf.arg.slot] = (h, w)
  if valid.op is not Ops.CONST or valid.val is not True:
    return buf.index(cidx.src[1].valid(valid), cidx.src[0].valid(valid))
@@ -67,7 +67,7 @@ def jit_lower(linear:UOp, held_bufs:set[UOp], input_uops:list[UOp]) -> UOp:
  if VIZ: graph_rewrite(linear, PatternMatcher([]), name="View captured linear")

  # parametrize input buffers: map each input buffer UOp to a PARAM with the correct slot index
-  linear = linear.substitute({u: UOp.param(i, u.dtype, u.max_numel(), u.device) for i,u in enumerate(input_uops)}, walk=True)
+  linear = linear.substitute({u: UOp.param(i, u.dtype, u.shape, u.device) for i,u in enumerate(input_uops)}, walk=True)
  linear = memory_plan_rewrite(linear, held_bufs)
  linear = compile_linear(linear, beam=getenv("JITBEAM", BEAM.value))
  if JIT < 2: linear = graph_split_rewrite(linear, max_batch_size=JIT_BATCH_SIZE.value)
@@ -145,7 +145,6 @@ def main():
  parser.add_argument("--serve", nargs='?', type=int, const=8000, metavar="PORT", help="Run OpenAI compatible API (optional port, default 8000)")
  parser.add_argument("--warmup", action="store_true", help="warmup the JIT")
  parser.add_argument("--benchmark", nargs='?', type=int, const=20, metavar="COUNT", help="Benchmark tok/s (optional count, default 20)")
-  parser.add_argument("--no_chat_template", action="store_true", help="Don't use the model's chat template, always use the fallback template")
  args = parser.parse_args()

  # load the model
@@ -161,7 +160,7 @@ def main():

  # use the model's chat template if jinja2 is available (enables model-specific formatting)
  template: jinja2.Template|FallbackTemplate = FallbackTemplate(tok)
-  if not args.no_chat_template and (ct := kv.get('tokenizer.chat_template')) is not None:
+  if (ct := kv.get('tokenizer.chat_template')) is not None:
    try:
      import jinja2
      env = jinja2.Environment()
@@ -31,10 +31,6 @@ def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]:
    return (None,) + k.arg.grad_fxn(on_dev(ctx, 0), k)
  assert fxn.op is Ops.TUPLE, f"expected TUPLE body for gradient, got {fxn.op}"
  params = {x.arg.slot:x for x in fxn.toposort(enter_calls=False) if x.op == Ops.PARAM}
-  # grads are collected at the flat param storage: reshape to each arg's view (max view shrunk to symbolic)
-  def shaped_grad(grad:UOp, i:int) -> UOp:
-    a = args[i]
-    return grad.view_as(a.shard_shape, a.axis) if a.axis is not None and isinstance(a.device, tuple) else grad.view_as(a._shape)
  grad_args = ctx.src
  root_grad = UOp(Ops.TUPLE, src=tuple(UOp(Ops.NOOP) if g.op is Ops.NOOP else
    g if g.device is None else g.param_like(len(args)+i) for i,g in enumerate(grad_args)))
@@ -43,7 +39,7 @@ def call_gradient(ctx:UOp, k:UOp, needed:set[int]) -> tuple[UOp|None, ...]:
  fwd_subs = {src: src.param_like(len(args)+len(grad_args)+i) for i, src in enumerate(fxn.src)} if k.arg.precompile else {}
  fwd_outs = tuple(k.gettuple(i) for i in range(len(fxn.src))) if k.arg.precompile else ()
  # collect needed gradient bodies, compact unused params, create a single backward CALL
-  grad_bodies = [(i, shaped_grad(grads[p], i)) for i in needed if (p:=params.get(i)) is not None and p in grads]
+  grad_bodies = [(i, grads[p]) for i in needed if (p:=params.get(i)) is not None and p in grads]
  bwd_body = UOp.maketuple(*(gb for _, gb in grad_bodies)).substitute(fwd_subs, walk=True)
  bwd_body = renumber_invalid_outputs(bwd_body)
  bwd_body, compact_args = _compact_params(bwd_body, (*args, *grad_args, *fwd_outs))
@@ -1067,7 +1067,7 @@ class OpMixin(ElementwiseMixin, ReduceMixin):
      reshape[i] = expand[i] = size[i]
      if mode == "linear":
        arr = type(self).arange(size[i])
-        num, den = (arr*(in_sz-1), max(size[i]-1, 1)) if align_corners else ((arr*2+1)*in_sz - size[i], size[i]*2)
+        num, den = (arr*(in_sz-1), size[i]-1) if align_corners else ((arr*2+1)*in_sz - size[i], size[i]*2)
        num = num.clip(0, (in_sz-1)*den)
        low, high, perc = [y.reshape(reshape).expand(expand) for y in (num//den, (num+den-1)//den, (num % den).cast(dtypes.float32)/den)]
        x = x.gather(i, low).lerp(x.gather(i, high), perc)
@@ -4,7 +4,7 @@ import sys, struct, functools
 from typing import cast
 from tinygrad.dtype import dtypes, DType, truncate, AddrSpace
 from tinygrad.uop import FastEnum, auto, Ops, GroupOp
-from tinygrad.uop.ops import UOp, UPat, PatternMatcher, promo_dtype
+from tinygrad.uop.ops import UOp, UPat, PatternMatcher
 from tinygrad.renderer.isa import ISARenderer, IselContext, Register, PreRegAllocContext, greg
 from tinygrad.helpers import getenv, NUM_CPU_THREADS, unwrap, Target

@@ -145,14 +145,14 @@ extra_matcher = PatternMatcher([
  # float16 alus are done in float32
  (UPat(GroupOp.ALU, dtypes.float16, name="x"), lambda x: UOp(x.op,
   src=tuple(s.cast(dtypes.float) if s.dtype != dtypes.bool else s for s in x.src)).cast(x.dtype)),
-  (UPat(GroupOp.Comparison, src=[UPat(dtype=dtypes.float16), UPat()], name="x"),
-   lambda x: UOp(x.op, src=tuple(s.cast(dtypes.float32) for s in x.src)).cast(x.dtype)),
+  (UPat(GroupOp.Comparison, src=(UPat.var("a", dtypes.float16), UPat.var("b")), name="x"),
+   lambda x,a,b: UOp(x.op, src=(a.cast(dtypes.float32), b.cast(dtypes.float32))).cast(x.dtype)),
  # no cmpne for packed ints, y != x => !(y==x)
  (UPat(Ops.CMPNE, src=(UPat.var("y", dtypes.ints), UPat.var("x")), name="cmp"),
   lambda y,x,cmp: UOp(Ops.CMPEQ, src=(y,x))^True if y.max_numel() > 1 else None),
-  # a float WHERE blends at the width of its value, so it needs a comparison at that width to make the mask
+  # float WHERE needs a mask unless its comparison already has a float operand
  (UPat.var("m", dtypes.bool).where(UPat.var("a", dtypes.floats+(dtypes.weakfloat,)), UPat.var("b")).named("w"),
-   lambda m,a,b,w: m.cast(w.dtype).ne(0).where(a, b) if w.dtype in dtypes.floats and promo_dtype(m.src) is not w.dtype else None),
+   lambda m,a,b,w: m.cast(w.dtype).ne(0).where(a, b) if w.dtype in dtypes.floats and not dtypes.is_float(m.src[0].dtype) else None),
  # rewrite -x -> 0 - x
  (UPat(Ops.NEG, name="x"), lambda x: UOp(Ops.SUB, src=(x.const_like(0),) + x.src)),
  # TODO: add support for mod, requires support for accessing the 2nd+ reg of a multi output instruction
@@ -282,7 +282,7 @@ def shift(x:UOp, op:X86Ops) -> UOp:
 # it is materialized as an immediate so the address stays correct if the base register is ever spilled and refilled
 def fold_address(x:UOp) -> tuple[UOp, UOp, UOp, UOp]:
  def _disp(v:int) -> UOp: return imm(dtypes.int32 if abs(v) > dtypes.int8.max else dtypes.int8, v)
-  def _cast(v:UOp) -> UOp: return v.cast(dtypes.int64) if v.vmin < 0 else v.cast(dtypes.uint32) if v.dtype.itemsize < 4 else v
+  def _cast(v:UOp) -> UOp: return v.cast(dtypes.int64) if v.vmin < 0 else v
  if x.op not in {Ops.INDEX, Ops.SHRINK}: return (x, UOp(Ops.NOOP), _disp(0), imm(dtypes.uint8, x.dtype.itemsize))
  base, idx = x.src[0], x.src[1]
  # buffers are indexed by element, everything else (the stack pointer) by byte
@@ -17,7 +17,6 @@ ggml_common_src = "https://raw.githubusercontent.com/ggml-org/ggml/d4fcfe88a8bcf
 cudart_src = "https://developer.download.nvidia.com/compute/cuda/redist/cuda_cudart/linux-x86_64/cuda_cudart-linux-x86_64-12.0.146-archive.tar.xz"
 nvrtc_src = "https://developer.download.nvidia.com/compute/cuda/redist/cuda_nvrtc/linux-x86_64/cuda_nvrtc-linux-x86_64-12.0.140-archive.tar.xz"
 opencl_src = "https://github.com/KhronosGroup/OpenCL-Headers/archive/2e30669d48718fd460f085b4b35b160dad51ce9d.tar.gz"
-comgr_2_src = "https://repo.radeon.com/rocm/apt/6.2/pool/main/c/comgr/comgr_2.8.0.60200-66~24.04_amd64.deb"
 macossdk = "/var/db/xcode_select_link/Platforms/MacOSX.platform/Developer/SDKs/MacOSX.sdk"

 llvm_lib = (
@@ -54,8 +53,6 @@ def load(name, files, **kwargs):
    if srcs: td.cleanup()
  return importlib.import_module(f"{path}.{name.replace('/', '.')}")

-def _extract_deb(path:str): subprocess.run("ar x *.deb && tar xf data.tar.*", cwd=path, shell=True, check=True)
-
 def __getattr__(nm):
  match nm:
    case "libc":
@@ -105,11 +102,14 @@ def __getattr__(nm):
    # this defines all syscall numbers. should probably unify linux autogen?
    case "io_uring":
      return load("io_uring", ["{}/liburing.h", "{}/usr/include/linux/io_uring.h", "{}/usr/include/asm-generic/unistd.h"],
-                  args=["-I{}/usr/include"], srcs=[linux_headers_deb, liburing_src], rules=[('__NR', 'NR')], preprocess=_extract_deb)
+                  args=["-I{}/usr/include"], srcs=[linux_headers_deb, liburing_src], rules=[('__NR', 'NR')],
+                  preprocess=lambda path: subprocess.run(f"ar x {linux_headers_deb.split('/')[-1]} && tar xf data.tar.xz", cwd=path, shell=True, check=True))
    case "llvm": return load("llvm", lambda: [system("llvm-config-20 --includedir")+"/llvm-c/**/*.h"], dll=llvm_lib,
                             args=lambda: system("llvm-config-20 --cflags").split(), recsym=True, prolog=["from tinygrad.helpers import WIN, OSX"])
-    case "pci": return load("pci", ["{}/usr/include/linux/pci_regs.h"], srcs=linux_headers_deb, preprocess=_extract_deb)
-    case "vfio": return load("vfio", ["{}/usr/include/linux/vfio.h"], args=["-I{}/usr/include"], srcs=linux_headers_deb, preprocess=_extract_deb)
+    case "pci": return load("pci", ["{}/usr/include/linux/pci_regs.h"], srcs=linux_headers_deb,
+                             preprocess=lambda path: subprocess.run(f"ar x {linux_headers_deb.split('/')[-1]} && tar xf data.tar.xz", cwd=path, shell=True, check=True))
+    case "vfio": return load("vfio", ["{}/usr/include/linux/vfio.h"], args=["-I{}/usr/include"], srcs=linux_headers_deb,
+                             preprocess=lambda path: subprocess.run(f"ar x {linux_headers_deb.split('/')[-1]} && tar xf data.tar.xz", cwd=path, shell=True, check=True))
    # could add rule: WGPU_COMMA -> ','
    case "webgpu": return load("webgpu", [root/"extra/webgpu/webgpu.h"], dll=webgpu_lib,
                               prolog=["from tinygrad.helpers import WIN, OSX", "import sysconfig, os"])
@@ -119,10 +119,9 @@ def __getattr__(nm):
                            dll="os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libamdhip64.so'",
                            args=["-D__HIP_PLATFORM_AMD__", "-I/opt/rocm/include", "-x", "c++"], prolog=["import os"])
    case "comgr" | "comgr_3":
-      prefix = "{}/opt/rocm-6.2.0" if nm == "comgr" else "/opt/rocm"
-      return load(nm, [f"{prefix}/include/amd_comgr/amd_comgr.h"], dll="[os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libamd_comgr.so', 'amd_comgr']",
-                  args=["-D__HIP_PLATFORM_AMD__", f"-I{prefix}/include", "-x", "c++"], prolog=["import os"], srcs=comgr_2_src if nm == "comgr" else None,
-                  **({'preprocess':_extract_deb} if nm == "comgr" else {}))
+      return load("comgr_3" if nm == "comgr_3" else "comgr", ["/opt/rocm/include/amd_comgr/amd_comgr.h"],
+                  dll= "[os.getenv('ROCM_PATH', '/opt/rocm')+'/lib/libamd_comgr.so', 'amd_comgr']",
+                  args=["-D__HIP_PLATFORM_AMD__", "-I/opt/rocm/include", "-x", "c++"], prolog=["import os"])
    case "hsa": return load("hsa", [*[f"{{}}/projects/rocr-runtime/runtime/hsa-runtime/core/inc/{s}.h" for s in ["registers"]],
                                    *[f"{{}}/projects/rocr-runtime/runtime/hsa-runtime/inc/{s}.h" for s in [
                                        "hsa", "hsa_ext_amd", "amd_hsa_signal", "amd_hsa_queue", "amd_hsa_kernel_code",
@@ -185,7 +184,8 @@ def __getattr__(nm):
      return load("mlx5", [root/"extra/mlx_driver/mlx5.h", f"{kh}/mlx5_ifc.h"], srcs=linux_headers_kern_deb,
                  args=["-Du8=unsigned char", "-Du16=unsigned short", "-Du32=unsigned int", "-Du64=unsigned long long",
                        "-D__be16=unsigned short", "-D__be32=unsigned int", "-D__be64=unsigned long long", f"-I{kh}"],
-                  preprocess=_extract_deb)
+                  preprocess=lambda path: subprocess.run(f"ar x {linux_headers_kern_deb.split('/')[-1]} && tar xf data.tar.xz",
+                                                         cwd=path, shell=True, check=True))
    case "bnxt":
      kh = "{}/usr/src/linux-headers-6.18.9+deb14-common/include"
      return load("bnxt", [f"{kh}/linux/bnxt/hsi.h", *[f"{{}}/{s.split('/')[-1]}" for s in bnxt_src]],
@@ -197,5 +197,7 @@ def __getattr__(nm):
                            r"((cmdq|creq)_(base|init|add_gid|create_(cq|qp)|initialize_fw|modify_qp|query_version|register_mr)(_resp)?"
                            r"|cq_(base|req)|sq_(rdma_hdr|sge))$",
                            r"(BNXT|CMDQ|CREQ|CQ|SQ|DBC|PTU|RCFW|HWRM|VNIC|RING_ALLOC|STAT_CTX|CFA_L2_FILTER|PORT_PHY_CFG|FIRMWARE_FIRST"
-                            r"|FUNC_(QCAPS|QCFG|RESET|DRV_RGTR|BACKING_STORE))_"], preprocess=_extract_deb)
+                            r"|FUNC_(QCAPS|QCFG|RESET|DRV_RGTR|BACKING_STORE))_"],
+                  preprocess=lambda path: subprocess.run(f"ar x {linux_headers_kern_deb.split('/')[-1]} && tar xf data.tar.xz",
+                                                         cwd=path, shell=True, check=True))
    case _: raise AttributeError(f"no such autogen: {nm}")
@@ -26,24 +26,24 @@ MAX_ARGS, CMD_SIZE, RING_SLOTS, FUNCS = 63, 64, (16 << 10), (() if WIN else ('cl

 def signal_prog():
  val = UOp.param(1, dtypes.int, (), vmin_vmax=(0, dtypes.int.max), name="value", addrspace=AddrSpace.ALU)
-  return UOp.param(0, dtypes.uint32, 1)[0].store(val.cast(dtypes.uint32))
+  return UOp.param(0, dtypes.uint32, (1,))[0].store(val.cast(dtypes.uint32))

 def wait_prog():
  val = UOp.param(1, dtypes.int, (), vmin_vmax=(0, dtypes.int.max), name="value", addrspace=AddrSpace.ALU)
-  return (v:=UOp.param(0, dtypes.uint32, 1, volatile=True).after(l:=UOp.loop(0))[0].load()).end(l, v < val.cast(dtypes.uint32))
+  return (v:=UOp.param(0, dtypes.uint32, (1,), volatile=True).after(l:=UOp.loop(0))[0].load()).end(l, v < val.cast(dtypes.uint32))

 def timestamp_prog():
  if WIN: val = UOp.const(0, dtypes.uint64)
  else:
-    fn, ts = UOp.param(1, dtypes.uint64, 1), UOp.placeholder((2,), dtypes.uint64, slot=0, addrspace=AddrSpace.REG)
+    fn, ts = UOp.param(1, dtypes.uint64, (1,)), UOp.placeholder((2,), dtypes.uint64, slot=0, addrspace=AddrSpace.REG)
    call = fn[0].load().call(UOp.const(6 if OSX else 1, dtypes.int), ts[0], ret_dtype=dtypes.void) # clock_gettime(CLOCK_MONOTONIC, &ts)
    val = ts.after(call)[0].load() * 1_000_000_000 + ts.after(call)[1].load()
-  return UOp.param(0, dtypes.uint64, 1)[0].store(val)
+  return UOp.param(0, dtypes.uint64, (1,))[0].store(val)

 def worker_prog():
-  ring = UOp.param(0, dtypes.uint64, RING_SLOTS * CMD_SIZE, volatile=True)
-  wait, done = UOp.param(1, dtypes.uint64, 1, volatile=True), UOp.param(2, dtypes.uint64, 1, volatile=True)
-  sem, cur = UOp.param(3, dtypes.uint64, 1), UOp.range(2**64-1, 0, dtype=dtypes.uint64) # sem is unused on windows, it has to come last
+  ring = UOp.param(0, dtypes.uint64, (RING_SLOTS * CMD_SIZE,), volatile=True)
+  wait, done = UOp.param(1, dtypes.uint64, (1,), volatile=True), UOp.param(2, dtypes.uint64, (1,), volatile=True)
+  sem, cur = UOp.param(3, dtypes.uint64, (1,)), UOp.range(2**64-1, 0, dtype=dtypes.uint64) # sem is unused on windows, it has to come last

  # spin on windows, sem_wait to sleep on posix
  if WIN: ready = (v:=wait.after(lw:=UOp.loop(1), cur)[0].load()).end(lw, v <= cur)
@@ -20,7 +20,7 @@ BUFTYPE_BUF, BUFTYPE_TEX, BUFTYPE_IBO = 0, 1, 2
 def dcache_flush():
  from tinygrad.uop.ops import UOp, Ops, KernelInfo
  from tinygrad.codegen import to_program
-  buf, n = UOp.param(0, dtypes.uint8, 1), UOp.param(1, dtypes.int, shape=(), name="n", addrspace=AddrSpace.ALU)
+  buf, n = UOp.param(0, dtypes.uint8, shape=(1,)), UOp.param(1, dtypes.int, shape=(), name="n", addrspace=AddrSpace.ALU)
  i = UOp.range(n, 0, dtype=dtypes.int)
  flush = UOp(Ops.CUSTOM, src=(buf.index(i * 64),), arg=('__asm__ volatile("dc cvac, %0" :: "r"({0}) : "memory");', dtypes.void))
  sink = UOp.sink(flush.end(i), UOp(Ops.CUSTOM, arg=('__asm__ volatile("dsb sy" ::: "memory");', dtypes.void)),
@@ -139,7 +139,7 @@ def _get_enqueue_devs(call:UOp) -> Any|None:

 def copy_with_kernel(call:UOp, dst:UOp, src:UOp) -> UOp|None:
  if (devs:=_get_enqueue_devs(call)) is None or Device[(dev:=to_tuple(devs)[0])].has_copy_queue: return None
-  d, s = (UOp.param(i, dst.dtype, n:=dst.max_numel(), device=devs) for i in range(2))
+  d, s = (UOp.param(i, dst.dtype, (n:=dst.max_numel(),), device=devs) for i in range(2))
  ast = d.index(r:=UOp.range(n, 0)).store(s.index(r).load()).end(r).sink(arg=KernelInfo(name="copy"), tag=1)
  return call.replace(src=(to_program(ast, Device[dev].renderer), dst, src))

@@ -81,7 +81,7 @@ def create_schedule(sched_sink:UOp) -> UOp:
 from tinygrad.schedule.memory import memory_plan_rewrite
 from tinygrad.engine.realize import capturing, pm_flatten_linear
 from tinygrad.schedule.prepare import prepare_rangeify
-from tinygrad.schedule.rangeify2 import get_kernel_graph
+from tinygrad.schedule.rangeify import get_kernel_graph
 from tinygrad.helpers import CAPTURING
 from tinygrad.uop.ops import PatternMatcher, UPat, ParamArg
 from tinygrad.dtype import AddrSpace
@@ -94,7 +94,7 @@ pm_post_sched_cache = PatternMatcher([
  # only resolve buffer PARAMs (slot>=0); ALU/shape vars use slot=-1 and must not be swapped for call args
  (UPat(Ops.PARAM, name="x"), lambda ctx,x: ctx[1][x.arg.slot] if x.arg.slot >= 0 else None),
  # create new BUFFERs
-  (UPat(Ops.BUFFER, src=(), name="b"), lambda ctx,b:
+  (UPat(Ops.BUFFER, src=(UPat(),), name="b"), lambda ctx,b:
   create_new_buffer(ctx, b) if isinstance(b.arg, ParamArg) and b.addrspace is AddrSpace.GLOBAL else None),
 ])

@@ -280,8 +280,13 @@ def rewrite_into_function(call:UOp):
                           for i, s in enumerate(new_body.src)])
  return call.replace(src=(new_body,)+new_args)

+def param_to_multi(p:UOp):
+  if p.axis is None: return None
+  return UOp.param(p.arg.slot, p.dtype, p.shard_shape, p.device, p.arg.vmin_vmax, p.arg.multiple_of, p.arg.name, p.arg.addrspace).unshard(p.axis)
+
 # NOTE: this is the same pattern as unrolled ranges
 multi_pm = PatternMatcher([
+  (UPat(Ops.PARAM, name="p"), param_to_multi),
  (UPat(GroupOp.ALU, name="root", custom_early_reject=set([Ops.UNSHARD])), alu_multi),
  (UPat(Ops.REDUCE, src=(UPat(Ops.UNSHARD, name="multi"), ), name="root"), reduce_multi),
  (UPat(Ops.RESHAPE, src=(UPat(Ops.UNSHARD, name="multi"), UPat()), name="root"), reshape_multi),
@@ -1,7 +1,7 @@
 import itertools
 from tinygrad.dtype import dtypes, to_dtype
 from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp
-from tinygrad.uop.ops import graph_rewrite, rewrite_group, ParamArg, identity_element
+from tinygrad.uop.ops import graph_rewrite, rewrite_group, shape_to_shape_arg, ParamArg, identity_element
 from tinygrad.uop.movement import mop_cleanup
 from tinygrad.helpers import prod, getenv, all_int, DEBUG, SPLIT_REDUCEOP, OPENPILOT_HACKS, FLOAT16, argsort
 from tinygrad.schedule.indexing import apply_movement_op
@@ -100,19 +100,10 @@ def resolve_function(c:UOp, allow_param_mismatch=True) -> UOp|None:
    if [x.arg.slot for x in params] != list(range(len(params))): raise RuntimeError(f"params not in order: {[x.arg.slot for x in params]}")
    if len(params) != len(args): raise TypeError(f"expected {len(params)} args, got {len(args)}")

-  # params have a flat storage size in the arg, the logical shape is a view (RESHAPE/SHRINK/UNSHARD) on top of it.
-  # substitute args by their flat max-shaped storage view so the movement views on the params stay valid
-  def flat_storage(a:UOp) -> tuple[int, UOp]:  # returns (size, view of a as flat max-shaped storage)
-    shp = a.max_shard_shape if a.axis is not None and isinstance(a.device, tuple) else a.max_shape
-    return (n:=prod(shp)), a if a.shape == (n,) else a.pad_to(shp).reshape((n,))
  dict_map = {x:args[x.arg.slot] for x in params}
  for i, (p, a) in enumerate(dict_map.items()):
-    if p.arg.size is not None:
-      n, flat = flat_storage(a)
-      if p.arg.size != n: raise TypeError(f"arg {i} shape mismatch: expected size {p.arg.size}, got {a.shape}")
-      dict_map[p] = flat
-    elif a.shape != ():
-      raise TypeError(f"arg {i} shape mismatch: expected scalar, got {a.shape}")
+    if p.axis != a.axis: raise TypeError(f"arg {i} axis mismatch: expected {p.axis}, got {a.axis}")
+    if p.max_shape != a.max_shape: raise TypeError(f"arg {i} shape mismatch: expected {p.shape}, got {a.shape}")
    if p.dtype != a.dtype: raise TypeError(f"arg {i} dtype mismatch: expected {p.dtype}, got {a.dtype}")
  return c.src[0].substitute(dict_map, walk=True)

@@ -201,13 +192,13 @@ def convert_copy_to_store(ctx, copy:UOp, existing_buf:UOp|None=None):
    # if there's already a buffer, we just use it
    return existing_buf.flatten().store(input_src)
  # create the output buffer
-  buf = UOp(Ops.BUFFER, arg=ParamArg(next(ctx), copy.dtype, size=prod(input_src.max_shape), device=copy.device))
+  buf = UOp(Ops.BUFFER, src=(shape_to_shape_arg(input_src.max_shape),), arg=ParamArg(next(ctx), copy.dtype, device=copy.device))
  # reshape back to input
-  return buf.reshape(input_src.max_shape).after(buf.store(input_src)).reshape(copy.shape)
+  return buf.after(buf.store(input_src)).reshape(copy.shape)

 pm_copy_to_store = PatternMatcher([
  (UPat(name="existing_buf").store(UPat(Ops.COPY, name="copy")), convert_copy_to_store),
-  (UPat((Ops.COPY, Ops.CONTIGUOUS), name="copy"), convert_copy_to_store),
+  (UPat(Ops.COPY, name="copy"), convert_copy_to_store),
 ])

@rewrite_group(new_ctx=False)
@@ -2,7 +2,7 @@ from dataclasses import dataclass, field, replace
 from typing import cast
 import itertools
 from tinygrad.dtype import dtypes, AddrSpace, Invalid
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, KernelInfo, ParamArg
+from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, resolve, GroupOp, KernelInfo, ParamArg, shape_to_shape_arg
 from tinygrad.uop.ops import graph_rewrite, sint, AxisType, BottomUpGate, rewrite_group
 from tinygrad.uop.symbolic import symbolic
 from tinygrad.helpers import prod, dedup, DEBUG_RANGEIFY, VIZ, MAX_KERNEL_BUFFERS, SPEC
@@ -162,12 +162,6 @@ pm_no_indexing_calls = PatternMatcher([
  (UPat(Ops.CALL, name="u"), no_indexing_calls),
 ])

-# the kernel graph is what gets executed: no shape views left in it, the storage of a value is just the storage
-pm_no_views = PatternMatcher([
-  (UPat((Ops.RESHAPE, Ops.SHRINK), name="v", src=(UPat((Ops.AFTER, Ops.PARAM, Ops.UNSHARD, Ops.MSTACK, Ops.BUFFER)),), allow_any_len=True), lambda v:
-   v.src[0]),
-])
-
 DEVICE_MAX_BUFS = {"METAL": 31, "WEBGPU": 8, "CPU": 31} # TODO: get from device?
@dataclass
 class LimitBufsContext:
@@ -229,7 +223,7 @@ def bufferize_to_store(ctx:itertools.count, x:UOp, idx:UOp, allow_locals=True):

  # NOTE: the local BUFFER needs to be disambiguated here
  if x.arg.addrspace == AddrSpace.GLOBAL:
-    buf = UOp(Ops.BUFFER, arg=ParamArg(next(ctx), x.dtype, size=size, device=x.arg.device, addrspace=AddrSpace.GLOBAL))
+    buf = UOp(Ops.BUFFER, src=(shape_to_shape_arg((size,)),), arg=ParamArg(next(ctx), x.dtype, device=x.arg.device, addrspace=AddrSpace.GLOBAL))
    do_store = buf.index(idx).store(x.src[0]).end(*rngs)
    return buf.after(do_store)

@@ -288,7 +282,8 @@ class LocalAddBufferContext:
 def debuf(ctx:LocalAddBufferContext, buf:UOp):
  # Variables (ALU buffers with a value range) are scalar symbolic values, not real buffers: they become ALU params with no slot
  if buf.is_variable: return buf.replace(op=Ops.PARAM)
-  param = UOp(Ops.PARAM, arg=ParamArg(ctx.dg, buf.dtype, prod(buf.max_shape), addrspace=buf.addrspace, device=buf.device))
+  param = UOp(Ops.PARAM, src=(UOp.const(prod(buf.max_shape)),),
+              arg=ParamArg(ctx.dg, buf.dtype, addrspace=buf.addrspace, device=buf.device))
  ret = param.reshape(buf.max_shape)
  # if the buffer has symbolic shape, shrink the max-sized view to the actual shape
  if buf.max_shape != buf.shape: ret = ret.shrink(tuple((0, s) for s in buf.shape))
@@ -363,7 +358,7 @@ def split_store(x:UOp) -> UOp|None:
  ret = graph_rewrite(x, to_define_global+pm_flatten_range+rangeify_codegen, ctx=lctx, name="kernel split", bottom_up=True)

  # create the Kernel. NOTE: buffers can be on different devices here now, they are compiled to SDMA copies later by schedule
-  return ret.sink(arg=KernelInfo(opts_to_apply=lctx.opts)).call(*lctx.map.values())
+  return ret.sink(arg=KernelInfo(opts_to_apply=lctx.opts)).call(*[x.flatten() for x in lctx.map.values()])

 split_kernels = PatternMatcher([
  (UPat((Ops.STORE, Ops.END), name="x"), split_store),
@@ -388,7 +383,6 @@ def get_kernel_graph(tsink:UOp) -> UOp:
  tsink = graph_rewrite(tsink, pm_add_buffers+pm_add_param_range_tags, ctx=itertools.count(paramarg_start), bottom_up=True, name="stage to store")
  tsink = graph_rewrite(tsink, split_kernels, bottom_up=True, name="split kernels")
  tsink = graph_rewrite(tsink, pm_no_indexing_calls, name="remove indexing from call args")
-  tsink = graph_rewrite(tsink, pm_no_views, name="remove views from the kernel graph")

  if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Kernel Graph")
  if SPEC:
@@ -1,230 +0,0 @@
-from dataclasses import dataclass, field
-import itertools
-from tinygrad.dtype import AddrSpace, Invalid
-from tinygrad.uop.ops import PatternMatcher, UPat, Ops, UOp, GroupOp, KernelInfo
-from tinygrad.uop.ops import graph_rewrite, AxisType, rewrite_group, remove_all_tags, resolve
-from tinygrad.helpers import all_int, VIZ, SPEC, Context, panic
-from tinygrad.schedule.indexing import BufferizeOpts, apply_movement_op, broadcast_rngs
-
-# *** RANGE creation ***
-
-def rangeify_on_reduce(ctx, inp:UOp, red:UOp, idx:UOp|None=None):
-  if red.arg[1] == 0: return None
-  if idx is None and len(red.shape) > 0: return None
-  # TODO: is AxisType.REDUCE a real thing?
-  rngs = [UOp.range(s, next(ctx), AxisType.REDUCE) for s in inp.shape[:red.arg[1]]]
-  return inp.index(*rngs, *(idx.src[1:] if idx is not None else ())).reduce(*rngs, arg=(red.arg[0], 0))
-
-def rangeify_on_store(ctx, x:UOp):
-  if x.shape == (): return None
-  rngs = [UOp.range(s, next(ctx)) for s in x.shape]
-  return x.src[0].index(*rngs).store(x.src[1].index(*rngs)).end(*rngs)
-
-def rangeify_on_stage(ctx, x:UOp):
-  if x.src[0].shape == (): return None
-  # size 1 dims don't get ranges, they are reshaped out and back in
-  if all_int(x.shape) and 0 < len(sq := tuple(s for s in x.shape if s != 1)) < len(x.shape):
-    return rangeify_on_stage(ctx, x.src[0].reshape(sq).bufferize(arg=x.arg)).reshape(x.shape)
-  rngs = [UOp.range(s, next(ctx)) for s in x.shape]
-  return x.replace(src=(x.src[0].index(*rngs), *rngs))
-
-pm_range_creation = PatternMatcher([
-  # reduce/store are what creates ranges
-  (UPat(Ops.REDUCE, src=(UPat.var('inp'),), name="red").index(name="idx", allow_any_len=True), rangeify_on_reduce),
-  (UPat(Ops.REDUCE, src=(UPat.var('inp'),), name="red"), rangeify_on_reduce),
-  (UPat(Ops.STORE, name="x"), rangeify_on_store),
-  (UPat(Ops.STAGE, name="x"), rangeify_on_stage),
-])
-
-# *** RANGE migration ***
-
-# movement op on INDEX as a PatternMatcher
-def _mop_index(r:UOp, idx:UOp):
-  idxs = idx.src[1:]
-  if len(idxs) == len(r.shape):
-    ret = r.src[0].index(*apply_movement_op(r.op, r.src[0].shape, r.marg, idxs), dtype=idx.dtype, arg=idx.arg)
-    if r.op is Ops.PAD:
-      # NOTE: neither 0 or ret.const_like(0) is correct here.
-      # const_like breaks because it adds casts, and 0 is wrong if ret is a bool
-      invalid_value = UOp.const(ret.dtype.const(0))
-      # insert invalid_value for PAD with where
-      a = UOp.const(True)
-      for s in UOp.sink(*ret.src[1:]).simplify().src:
-        if s.is_invalid: return invalid_value
-        if s.op is Ops.WHERE and s.src[2].op is Ops.CONST and s.src[2].arg == Invalid: a = a & s.src[0]
-      ret = a.where(ret, invalid_value)
-    return ret
-  if r.op is Ops.RESHAPE:
-    src_prefix = len(r.src[0].shape) - len(r.shape[len(idxs):])
-    if src_prefix >= 0 and r.src[0].shape[src_prefix:] == r.shape[len(idxs):]:
-      if src_prefix == 0: return r.src[0] if r.src[0].dtype == idx.dtype else None
-      ret = r.src[0].index(*apply_movement_op(r.op, r.src[0].shape[:src_prefix], r.shape[:len(idxs)], idxs), dtype=idx.dtype, arg=idx.arg)
-      return ret if ret.shape == idx.shape else None
-
-# TODO: this should be in _mop_index
-def index_on_stack(stack:UOp, idx:UOp):
-  srcs = [s.index(*idx.src[2:]) for s in stack.src]
-  r0 = idx.src[1]
-  ret = srcs[-1]
-  for k in range(len(srcs)-2, -1, -1): ret = r0.eq(k).where(srcs[k], ret)
-  return ret
-
-pm_range_migration = PatternMatcher([
-  # STAGE on shape () is nothing
-  (UPat(Ops.STAGE, src=(UPat.var('x'),)), lambda x: x if x.shape == () else None),
-  # if INDEX is on STAGE with the same ranges, remove the pair
-  (UPat(Ops.STAGE, allow_any_len=True, name="s").index(allow_any_len=True, name="i"),
-   lambda s,i: s.src[0] if s.src[1:] == i.src[1:] else None),
-  # RESHAPE to scalar is INDEX
-  (UPat(Ops.RESHAPE, name="x"), lambda x: x.src[0].index(*[0 for _ in range(x.src[0].ndim)]) if x.shape == () else None),
-  # handle movement ops on INDEX
-  (UPat(GroupOp.Movement, name="r").index(name="idx", allow_any_len=True), _mop_index),
-  (UPat(Ops.STACK, name="stack").index(name="idx", allow_any_len=True), index_on_stack),
-  # move movement ops and INDEX after AFTER
-  (UPat(GroupOp.Movement|{Ops.INDEX}, name="r").after(name="a", allow_any_len=True),
-   lambda r,a: UOp(r.op, src=(a.replace(src=(r.src[0],)+a.src[1:]),)+r.src[1:], arg=r.arg)),
-  # block bitcast that changes shape
-  (UPat(Ops.BITCAST, name="b").index(allow_any_len=True),
-   lambda b: panic(RuntimeError, "shape changing bitcast not allowed in rangeify") if b.src[0].shape != b.shape else None),
-  # pass index through elementwise (with broadcasting support)
-  (UPat(GroupOp.Elementwise, name="b").index(name="idx", allow_any_len=True),
-   lambda b,idx: b.replace(src=tuple(s.index(*broadcast_rngs(b, s, idx.src[1:])) for s in b.src))),
-  # INDEX without src is nothing (must be at the bottom)
-  (UPat(Ops.INDEX, src=(UPat.var('x'),)), lambda x: x),
-])
-
-# *** split into kernels ***
-
-@dataclass
-class SplitCtx:
-  call_args:list = field(default_factory=list)
-  range_number:int = -1
-  addrspace:AddrSpace = AddrSpace.GLOBAL
-
-def _split_graph(ctx:SplitCtx, u:UOp) -> UOp|None:
-  if u.tag is not None: return None
-  if u.addrspace != ctx.addrspace: return None
-  # ALU params are scalar symbolic values, not buffers: they are already params and bind by name in the outer LINEAR
-  if u.op is Ops.PARAM and u.arg.addrspace is AddrSpace.ALU: return None
-  us = u.flatten() if u.addrspace == AddrSpace.GLOBAL else u
-  ctx.call_args.append(us)
-  return us.param_like(len(ctx.call_args)-1).rtag().reshape(u.shape)
-
-def _renumber_range(ctx:SplitCtx, u:UOp) -> UOp|None:
-  if u.tag is not None: return None
-  ctx.range_number += 1
-  return u.replace(arg=(ctx.range_number, u.arg[-1])).rtag()
-
-pm_split_graph = pm_range_migration+PatternMatcher([
-  (UPat((Ops.PARAM, Ops.AFTER, Ops.BUFFER, Ops.MSELECT, Ops.MSTACK), name="u"), _split_graph),
-  (UPat(Ops.RANGE, name="u"), _renumber_range),
-])
-
-def split_store(x:UOp) -> UOp:
-  ret = graph_rewrite(x, pm_split_graph, ctx:=SplitCtx(), name="split kernel", bottom_up=True)
-  # TODO: params and args should be able to be in any order
-  ctx.addrspace = AddrSpace.ALU
-  ret = graph_rewrite(ret, pm_split_graph, ctx, name="split kernel (vars)", bottom_up=True)
-  ret = graph_rewrite(ret, remove_all_tags, name="remove split tags", bottom_up=True)
-  return ret.sink(arg=KernelInfo()).call(*ctx.call_args)
-
-split_kernels = PatternMatcher([
-  (UPat((Ops.STORE, Ops.END), name="x"), split_store),
-])
-
-# cleanups
-
-def strip_zero_offset_shrink(x:UOp) -> UOp:
-  return x.src[0] if x.op is Ops.SHRINK and all(resolve(start == 0, False) for start,_ in x.marg) else x
-
-def no_indexing_calls(u:UOp):
-  new_srcs = []
-  for x in u.src:
-    if x.op is Ops.INDEX:
-      # sometimes if call srcs have children the call will get an INDEX. we remove it here.
-      # TODO: we should add safety checks here for contiguous
-      new_srcs.append(x.src[0])
-    elif x.op is Ops.SHRINK:
-      # SHRINK with offset 0 is fine
-      new_srcs.append(strip_zero_offset_shrink(x))
-    elif x.op is Ops.MSTACK:
-      new_srcs.append(x.replace(src=tuple(strip_zero_offset_shrink(s) for s in x.src)))
-    else:
-      # everything else we pass through
-      new_srcs.append(x)
-  return u.replace(src=tuple(new_srcs))
-
-pm_no_indexing_calls = PatternMatcher([
-  (UPat(Ops.CALL, name="u"), no_indexing_calls),
-])
-
-# the kernel graph is what gets executed: no shape views left in it, the storage of a value is just the storage
-pm_no_views = PatternMatcher([
-  (UPat((Ops.RESHAPE, Ops.SHRINK, Ops.INDEX), name="v"),lambda v: v.src[0]),
-])
-
-# *** rangeify ***
-
-debug_tag_factor = PatternMatcher([
-  (UPat(GroupOp.All, name="x"), lambda ctx,x: x.rtag(ctx[0][x] if x not in ctx[1] else 'REAL') if x.tag is None else None),
-])
-
-def remove_stage(ctx, x:UOp) -> UOp:
-  buf = UOp.new_buffer(x.arg.device, x.max_numel(), x.dtype, num=next(ctx))
-  return buf.after(buf.reshape(x.shape).index(*x.src[1:]).store(x.src[0]).end(*x.src[1:])).reshape(x.shape)
-
-pm_remove_stage = PatternMatcher([
-  (UPat(Ops.STAGE, name="x"), remove_stage),
-])
-
-@rewrite_group(new_ctx=False)
-def get_kernel_graph(tsink:UOp) -> UOp:
-  # add safe STAGEs to never duplicate compute
-  # we compute the number of times a buffer is consumed. if > 1, we realize
-  realize = {}
-  consumes = {tsink:0}
-  for u in reversed(tsink.toposort()):
-    assert u in consumes, f"{u.op} not in consumes"
-    if (u.op in GroupOp.ALU or u.op is Ops.REDUCE) and consumes[u] > 1 and u.device is not None:
-      # TODO: rename to stage
-      realize[u] = u.rtag(1).bufferize(arg=BufferizeOpts(device=u.device))
-      consumes[u] = 1
-    if u.op is Ops.STORE: consumes[u] = 1
-    if u.op is Ops.EXPAND: consumes[u] *= u.max_numel() // u.src[0].max_numel()
-    for i,s in enumerate(u.src):
-      if s not in consumes: consumes[s] = 0
-      if u.op is not Ops.STORE or i > 0:
-        consumes[s] += consumes[u]
-
-  if VIZ:
-    with Context(TRACK_MATCH_STATS=0): ctags = graph_rewrite(tsink, debug_tag_factor, ctx=(consumes, realize), bottom_up=True)
-    graph_rewrite(ctags, PatternMatcher([]), name="View Consumes")
-
-  # add stages
-  tsink = graph_rewrite(tsink.substitute(realize), remove_all_tags, name="untag")
-
-  # simple rangeify
-  tsink = graph_rewrite(tsink, pm_range_creation+pm_range_migration, ctx=itertools.count(0), bottom_up=True, name="simple rangeify")
-
-  if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Rangeify")
-
-  # ***** MERGING AND SPLITTING (should be totally optional) *****
-
-  # TODO
-
-  # ***** MERGING AND SPLITTING (should be totally optional) *****
-
-  if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Merged Rangeify")
-
-  next_buffer_num = itertools.count(1000)
-  tsink = graph_rewrite(tsink, pm_remove_stage, ctx=next_buffer_num, bottom_up=True, name="remove stage")
-  tsink = graph_rewrite(tsink, split_kernels, bottom_up=True, name="split kernels")
-  tsink = graph_rewrite(tsink, pm_no_indexing_calls, name="remove indexing from call args")
-  tsink = graph_rewrite(tsink, pm_no_views, name="remove views from the kernel graph")
-
-  if VIZ: graph_rewrite(tsink, PatternMatcher([]), name="View Kernel Graph")
-  if SPEC:
-    # validate the kernel graph
-    from tinygrad.uop.spec import type_verify, spec_kernel_graph
-    type_verify(tsink, spec_kernel_graph, enter_calls=False)
-  return tsink
@@ -198,7 +198,9 @@ def finalize_after(ctx:AllocCtx, x:UOp):

 def replace_input_buffer(ctx:AllocCtx, b:UOp):
  ctx.replacements.append(b)
-  return b.param_like(len(ctx.replacements)-1)
+  if b.is_bound_var or b.is_variable: return b.param_like(len(ctx.replacements)-1)
+  return UOp.param(len(ctx.replacements)-1, b.dtype, b.shape, b.device,
+                   addrspace=b.addrspace if b.addrspace is not None else AddrSpace.GLOBAL)

 def replace_input_view(ctx:AllocCtx, b:UOp): return replace_input_buffer(ctx, b) if b in ctx.views else None

@@ -209,7 +211,7 @@ pm_finalize_call = PatternMatcher([

 pm_replace_buf = PatternMatcher([
  # replace BUFFER with PARAM for cache key normalization
-  (UPat(Ops.BUFFER, src=(), name="b"), lambda ctx,b:
+  (UPat(Ops.BUFFER, src=(UPat(),), name="b"), lambda ctx,b:
   replace_input_buffer(ctx, b) if isinstance(b.arg, ParamArg) and b.addrspace is AddrSpace.GLOBAL else None),
  # replace SHRINK with PARAM
  (UPat(Ops.SHRINK, src=(UPat(Ops.BUFFER),), name="b", allow_any_len=True), replace_input_view),
@@ -379,7 +381,7 @@ class Tensor(RandMixin):
  # ***** data handlers ****

  def as_param(self, slot:int):
-    return Tensor(self.uop.param_like(slot))
+    return Tensor(UOp.param(slot, self.dtype, self.uop.shard_shape, self.device, axis=self.uop.axis))

  def call(self, *lst:Tensor, fxn:Tensor|UOp, grad_fxn:Callable|None=None) -> Tensor:
    fret = fxn._uop.call(*[t.uop for t in (self,)+lst], grad_fxn=grad_fxn)
@@ -23,21 +23,17 @@ class AxisType(Enum):
 class ParamArg:
  slot: int
  dtype: DType
-  # number of elements in the buffer. always a concrete int (never symbolic), None for scalars (shape ())
-  size: int|None = None
  vmin_vmax: tuple[PyConst, PyConst]|None = None
  multiple_of: int|None = None
  name: str|None = None
  addrspace: AddrSpace|None = AddrSpace.GLOBAL
+  axis: int|None = None
  device: str|tuple[str, ...]|None = None
  volatile: bool = False
-  # (h, w) if this is an image2d buffer, then size == h*w*4
-  image: tuple[int, int]|None = None
  def __repr__(self):
-    fields = (("vmin_vmax", None), ("multiple_of", None), ("name", None), ("addrspace", AddrSpace.GLOBAL), ("device", None),
-              ("volatile", False), ("image", None))
-    args = [repr(self.slot), repr(self.dtype)] + ([repr(self.size)] if self.size is not None else []) + \
-      [f"{k}={v!r}" for k,default in fields if (v:=getattr(self, k)) != default]
+    fields = (("vmin_vmax", None), ("multiple_of", None), ("name", None), ("addrspace", AddrSpace.GLOBAL), ("axis", None), ("device", None),
+              ("volatile", False))
+    args = [repr(self.slot), repr(self.dtype)] + [f"{k}={v!r}" for k,default in fields if (v:=getattr(self, k)) != default]
    return f"ParamArg({', '.join(args)})"
 axis_letters = {AxisType.DEVICE: "d", AxisType.GLOBAL: "g", AxisType.THREAD: "t", AxisType.LOCAL: "l", AxisType.WARP: "w", AxisType.WEAK: "L",
                AxisType.LOOP: "L", AxisType.UPCAST: "u", AxisType.GROUP_REDUCE: "G", AxisType.REDUCE: "R", AxisType.UNROLL: "r"}
@@ -385,10 +381,9 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
      case Ops.GETADDR: return ()
      case Ops.RANGE | Ops.SPECIAL: return ()
      case Ops.BINARY: return (len(self.arg),)
-      case Ops.BUFFER | Ops.PARAM:
-        # PARAM/BUFFER don't have a shape input, they have a size in the arg: int gives shape (size,), None gives ()
-        if (img:=self.arg.image) is not None: return (img[0], img[1], 4)
-        return () if self.arg.size is None else (self.arg.size,)
+      case Ops.BUFFER:
+        if len(self.src): return self.src[0].as_shape
+        return ()
      case Ops.CUSTOM | Ops.CUSTOMI:
        if self.dtype is dtypes.void: return None
        input_shapes = [x._shape for x in self.src if x._shape is not None]
@@ -399,6 +394,10 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
        # STAGE adds the existing shape to the front, opposite of INDEX
        return tuple([int(r.vmax+1) for r in self.src[1:]])+self.src[0].shape

+      # param has shape as the only arg
+      case Ops.PARAM:
+        return self.src[0].as_shape
+
      # wmma output shape = accumulator shape (src[2])
      case Ops.WMMA:
        wmma_b = _broadcast_shape(self.src[0].shape[:-1], self.src[1].shape[:-1], self.src[2].shape[:-1])
@@ -709,7 +708,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
    if self.op is Ops.GETTUPLE:
      in_tuple = self.src[0].src[0] if self.src[0].op is Ops.FUNCTION else self.src[0]
      return in_tuple.src[self.arg].axis if in_tuple.op is Ops.TUPLE else None
-    if self.op is Ops.PARAM: return None
+    if self.op is Ops.PARAM: return self.arg.axis
    # NOTE: they all have to share an axis, we always choose [-1]. src axes are right-aligned into the output shape
    if self.op in GroupOp.ALU.union({Ops.STACK}):
      return axes[-1] if (axes := dedup([x.axis+len(self.shape)-len(x.shape) for x in self.src if x.axis is not None])) else None
@@ -822,9 +821,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
  @staticmethod
  def new_buffer(device:str|tuple[str, ...], size:int, dtype:DType, num=None):
    if dtype in dtypes.weaks: raise RuntimeError(f"cannot create storage for weak dtype {dtype}")
-    assert isinstance(size, int), f"new_buffer size must be a concrete int, got {size}"
    slot = next(UOp.unique_num) if num is None else num
-    return UOp(Ops.BUFFER, arg=ParamArg(slot, dtype, size=size, device=device))
+    return UOp(Ops.BUFFER, src=(shape_to_shape_arg((size,)),), arg=ParamArg(slot, dtype, device=device))
  @staticmethod
  def from_buffer(opaque:Buffer, device:str|tuple[str, ...]|None=None):
    if (uop:=UOp.new_buffer(device or opaque.device, opaque.size, opaque.dtype, num=-id(opaque))) not in buffers: buffers[uop] = opaque.ref(1)
@@ -980,7 +978,7 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
    # a Variable is a 0-d BUFFER in the ALU addrspace; binding it is storing a CONST into it
    # param=True creates the kernel-side form directly: an ALU PARAM (what the BUFFER becomes inside kernels)
    arg = ParamArg(-1, dtype, name=name, vmin_vmax=(min_val, max_val), multiple_of=multiple_of, addrspace=AddrSpace.ALU)
-    return UOp(Ops.PARAM if param else Ops.BUFFER, arg=arg)
+    return UOp(Ops.PARAM if param else Ops.BUFFER, src=(shape_to_shape_arg(()),), arg=arg)
  @property
  def is_variable(self) -> bool:
    # a Variable is a 0-d BUFFER in the ALU addrspace that carries a value range (it becomes a PARAM inside kernels)
@@ -1150,11 +1148,11 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
    dtype = strong_dtype(dtype)  # storage is never weak: a placeholder commits the width of what's put in it
    if slot is None: slot = next(UOp.unique_num)
    if addrspace is AddrSpace.GLOBAL:
-      ret = UOp(Ops.PARAM, arg=ParamArg(slot, dtype, size=prod(shape), addrspace=addrspace, device=device, volatile=volatile))
+      ret = UOp(Ops.PARAM, src=(shape_to_shape_arg((prod(shape),)),), arg=ParamArg(slot, dtype, addrspace=addrspace, device=device,volatile=volatile))
    else:
      assert addrspace in (AddrSpace.LOCAL, AddrSpace.REG)
      assert device is None, "LOCAL and REG placeholders cannot have a device"
-      ret = UOp(Ops.BUFFER, arg=ParamArg(slot, dtype, size=prod(shape), addrspace=addrspace))
+      ret = UOp(Ops.BUFFER, src=(shape_to_shape_arg((prod(shape),)),), arg=ParamArg(slot, dtype, addrspace=addrspace))
    if tag is not None: ret = ret.rtag(tag)
    if len(shape) > 1: ret = ret.reshape(shape)
    return ret
@@ -1168,35 +1166,20 @@ class UOp(RandMixin, metaclass=UOpMetaClass):

  # TODO: this should replace placeholder
  @staticmethod
-  def param(slot:int, dtype:DType, shape:tuple[sint, ...]|sint|None=None, device=None, vmin_vmax:tuple[PyConst, PyConst]|None=None,
-            multiple_of:int|None=None, name=None, addrspace=AddrSpace.GLOBAL, volatile:bool=False):
-    """create a PARAM: a single sint or 1-d shape gives a flat param of that size, a None shape gives a scalar param.
-    the arg only stores the concrete max size (never symbolic): a multi-dim shape is a RESHAPE on top of the flat param,
-    a symbolic shape is a max-size param shrunk to the real shape"""
+  def param(slot:int, dtype:DType, shape:tuple[sint, ...]|None=None, device=None, vmin_vmax:tuple[PyConst, PyConst]|None=None,
+            multiple_of:int|None=None, name=None, addrspace=AddrSpace.GLOBAL, axis:int|None=None, volatile:bool=False):
    if dtype in dtypes.weaks: raise RuntimeError(f"cannot create param for weak dtype {dtype}")
-    if isinstance(shape, (int, UOp)): shape = (shape,)
-    if shape is None or len(shape) == 0:
-      return UOp(Ops.PARAM, arg=ParamArg(slot, dtype, None, vmin_vmax, multiple_of, name, addrspace, device, volatile))
-    max_shape = to_max_shape(shape)
-    ret = UOp(Ops.PARAM, arg=ParamArg(slot, dtype, prod(max_shape), vmin_vmax, multiple_of, name, addrspace, device, volatile))
-    return ret.view_as(shape)
+    if shape is not None and axis is not None and isinstance(device, tuple):
+      shape = tuple(s*len(device) if i == axis else s for i,s in enumerate(shape))
+    src: tuple[UOp, ...] = (UOp(Ops.NOOP) if shape is None else shape_to_shape_arg(shape),)
+    return UOp(Ops.PARAM, src=src, arg=ParamArg(slot, dtype, vmin_vmax, multiple_of, name, addrspace, axis, device, volatile))
  def param_like(self, slot:int):
    # Variables become ALU params in the call body; the stored value (if bound) stays in the call args
    if self.is_bound_var or self.is_variable:
      b = self.src[0] if self.op is Ops.AFTER else self
-      return UOp(Ops.PARAM, arg=replace(b.arg, slot=slot, name=f"p{slot}"))
+      return UOp(Ops.PARAM, src=b.src, arg=replace(b.arg, slot=slot, name=f"p{slot}"))
    addrspace = self.addrspace if self.addrspace is not None else AddrSpace.GLOBAL
-    # multi-device values become a per-shard sized param wrapped in UNSHARD: the sharding lives in the graph, not the arg
-    if self.axis is not None and isinstance(self.device, tuple):
-      return UOp(Ops.PARAM, arg=ParamArg(slot, self.dtype, prod(to_max_shape(self.shard_shape)),
-                                         addrspace=addrspace, device=self.device)).view_as(self.shard_shape, self.axis)
-    return UOp.param(slot, self.dtype, self._shape, self.device, addrspace=addrspace)
-  def view_as(self:UOp, shape:tuple[sint, ...], axis:int|None=None) -> UOp:
-    """view flat storage as the given (possibly symbolic) shape, optionally sharded on axis, the UNSHARD gives back the multiplied shape"""
-    max_shape = to_max_shape(shape)
-    ret = self.reshape(max_shape) if len(shape) > 1 else self
-    if tuple(max_shape) != tuple(shape): ret = ret.shrink_to(shape)
-    return ret if axis is None else ret.unshard(axis)
+    return UOp.param(slot, self.dtype, self.shard_shape if self.axis is not None else self._shape, self.device, addrspace=addrspace, axis=self.axis)

  @staticmethod
  def custom_function(name:str, *src:UOp) -> UOp: return UOp(Ops.CUSTOM_FUNCTION, src=src, arg=name)
@@ -1215,8 +1198,8 @@ class UOp(RandMixin, metaclass=UOpMetaClass):
    body = self if self.op is Ops.TUPLE else UOp.maketuple(self)
    return UOp(Ops.FUNCTION, src=(body,)+srcs, arg=CallInfo(grad_fxn, name, precompile, precompile_backward, aux))
  def custom_kernel(*srcs:UOp, fxn:Callable, grad_fxn:Callable|None=None) -> list[UOp]:
-    placeholders = [UOp.placeholder_like(s, slot=i) for i,s in enumerate(srcs)]
-    kernel = fxn(*placeholders).call(*srcs, grad_fxn=grad_fxn)
+    placeholders = [UOp.placeholder_like(s.flatten(), slot=i).reshape(s.shape) for i,s in enumerate(srcs)]
+    kernel = fxn(*placeholders).call(*[s.flatten() for s in srcs], grad_fxn=grad_fxn)
    return [s.after(kernel) for s in srcs]

  def to_elf(self) -> TinyELF:
@@ -84,7 +84,7 @@ pm_pyrender_extra = PatternMatcher([
  (UPat(Ops.CONST, src=(), name="x"), lambda x: f"UOp.const({x.val})"),
  (UPat((Ops.CAST, Ops.BITCAST), name="x"), lambda ctx,x: f"{ctx[x.src[0]]}.{x.op.name.lower()}({x.dtype})" if x.dtype != x.src[0].dtype else None),
  (UPat(Ops.SPECIAL, src=(UPat(Ops.CONST),), name="x"), lambda x: f"UOp.special({x.src[0].val}, {repr(x.arg)})"),
-  (UPat(Ops.BUFFER, src=(), name="x"), lambda x:
+  (UPat(Ops.BUFFER, src=(UPat(),), name="x"), lambda x:
    f"UOp.new_buffer({repr(x.arg.device)}, {x.max_numel()}, {x.dtype}, {x.arg.slot})"
    if isinstance(x.arg, ParamArg) and x.addrspace is AddrSpace.GLOBAL else None),
  (UPat(Ops.COPY, src=(UPat(name="x"),), name="copy"), lambda ctx,x,copy: f"{ctx[x]}.copy_to_device({repr(copy.arg)})"),
@@ -87,9 +87,10 @@ spec_shared = PatternMatcher([
  # a loop-ended END requires a trailing bool condition for the backedge (loop again while true)
  (UPat(Ops.END, src=(UPat(), UPat(Ops.RANGE, dtypes.void), UPat(dtype=dtypes.bool))), lambda: True),

-  # PARAM/BUFFER have a size in the arg, no shape input
-  (UPat(Ops.PARAM, src=(), name="x"), lambda x: isinstance(x.arg, ParamArg)),
-  (UPat(Ops.BUFFER, src=(), name="x"), lambda x: isinstance(x.arg, ParamArg) and x.addrspace in (AddrSpace.REG, AddrSpace.LOCAL)),
+  # PARAM
+  (UPat(Ops.PARAM, name="x"), lambda x: isinstance(x.arg, ParamArg)),
+  (UPat(Ops.BUFFER, src=(UPat(),), name="x"), lambda x:
+   isinstance(x.arg, ParamArg) and x.addrspace in (AddrSpace.REG, AddrSpace.LOCAL)),

  # GROUP of stores (or groups, or NOOPs)
  (UPat(Ops.GROUP, dtypes.void, src=UPat((Ops.GROUP, Ops.STORE, Ops.NOOP, Ops.INS, Ops.END))), lambda: True),
@@ -140,12 +141,12 @@ spec_tensor = PatternMatcher([
   lambda u: dtypes.is_float(u.dtype) or u.src[0].base.is_invalid),

  # BUFFER
-  (UPat(Ops.BUFFER, src=(), name="buf"), lambda buf:
-   (isinstance(buf.dtype, DType) and isinstance(buf.arg.size, int) and is_device(buf.arg.device))
+  (UPat(Ops.BUFFER, src=(UPat(),), name="buf"), lambda buf:
+   (isinstance(buf.dtype, DType) and matches_dtype(buf.src[0], dtypes.weakint) and is_device(buf.arg.device))
   if isinstance(buf.arg, ParamArg) and buf.addrspace is AddrSpace.GLOBAL else None),

  # a Variable is a 0-d ALU BUFFER with a value range and no device
-  (UPat(Ops.BUFFER, src=(), name="buf"), lambda buf: buf.arg.device is None if buf.is_variable else None),
+  (UPat(Ops.BUFFER, src=(UPat(),), name="buf"), lambda buf: buf.arg.device is None if buf.is_variable else None),

  # custom function
  (UPat(Ops.CUSTOM_FUNCTION, name="x"), lambda x: isinstance(x.arg, str)),
@@ -259,15 +260,17 @@ spec_kernel_graph = PatternMatcher([
  (UPat(Ops.STACK, name="s"), lambda s: all(x.op in (Ops.CONST, Ops.PARAM) or x.is_variable or x.is_bound_var for x in s.src) or None),
  # linear for more kernels (TODO: we should enter non sink calls)
  #(UPat(Ops.LINEAR), lambda: True),
-  # param is outside buffer, buffer is local buffer. params have a size in the arg, no shape input
-  (UPat(Ops.PARAM, src=(), name="x"), lambda x: isinstance(x.arg, ParamArg)),
+  # param is outside buffer, buffer is local buffer
+  (UPat(Ops.PARAM, name="x"), lambda x: isinstance(x.arg, ParamArg)),
  (UPat(Ops.BUFFER, name="x"), lambda x: isinstance(x.arg, ParamArg) and x.addrspace in (AddrSpace.GLOBAL, AddrSpace.ALU)),
-  (UPat(Ops.BITCAST), lambda: True),
+  # RESHAPE/BITCAST are NOOPs in the kernel graph (do we need them?)
+  (UPat((Ops.RESHAPE, Ops.BITCAST)), lambda: True),
  # mstack/mselect
  (UPat(Ops.MSTACK, name="x"), lambda x: all(isinstance(s.device, str) for s in x.src) or (all_same(x.src) and x.src[0].device is None)),
  (UPat(Ops.MSELECT, name="x"), lambda x: isinstance(x.src[0].device, tuple) and x.arg < len(x.src[0].device)),
-  # all calls are on various sinks
-  (UPat(Ops.CALL, src=(UPat((Ops.SINK, Ops.LINEAR, Ops.PROGRAM, Ops.CUSTOM_FUNCTION)),), allow_any_len=True), lambda: True),
+  # CALL (for kernels, all input args must be 0 or 1-d)
+  (UPat(Ops.CALL, src=(UPat((Ops.SINK, Ops.LINEAR, Ops.PROGRAM, Ops.CUSTOM_FUNCTION)),), allow_any_len=True, name="c"),
+   lambda c: all(u.ndim <= 1 for u in c.src[1:])),
  # after on PARAM or AFTER
  (UPat(Ops.AFTER, src=(UPat(GroupOp.Movement.union({Ops.PARAM, Ops.AFTER, Ops.BUFFER, Ops.MSTACK, Ops.MSELECT, Ops.BITCAST, Ops.RESHAPE})),),
        allow_any_len=True, name="x"), lambda x: matches_dtype(x.src[0], x.dtype)),
Author	SHA1	Message	Date
geohot	6234b42cf7	low rank call in kernel graph	2026-08-27 10:34:54 -07:00
geohot	ae9da86538	fix custom_kernel wrapper	2026-08-27 10:18:37 -07:00
geohot	df6df141ec	call has rank <= 1	2026-08-27 10:15:35 -07:00