diff --git a/tests/nki/test_nki_kernel_idioms.py b/tests/nki/test_nki_kernel_idioms.py
new file mode 100644
index 00000000..494fa78d
--- /dev/null
+++ b/tests/nki/test_nki_kernel_idioms.py
@@ -0,0 +1,140 @@
+"""Kernels written in the style of the official NKI samples, run by TileLens and by
+NKI's own CPU simulator (``nki.simulate_kernel``); both must give the same outputs."""
+
+import math
+
+import numpy as np
+import pytest
+
+try:
+    from neuronxcc import nki
+    import neuronxcc.nki.language as nl
+except ModuleNotFoundError:
+    pytest.skip(
+        "NeuronX dependencies are missing. Install tilelens[nki] to run these tests.",
+        allow_module_level=True,
+    )
+
+import tilelens
+from tilelens.clients import Tracer
+
+pytestmark = pytest.mark.nki  # only run at "pytest -m nki"
+
+
+def softmax_kernel(a_tensor):
+    # The output is allocated in shared HBM and returned; rows are tiled by 128 and
+    # the tail is masked with a [128, 1] mask that broadcasts over the tile.
+    out_tensor = nl.ndarray(a_tensor.shape, dtype=a_tensor.dtype, buffer=nl.shared_hbm)
+    ix = nl.arange(128)[:, None]
+    iy = nl.arange(a_tensor.shape[1])[None, :]
+    num_rows = a_tensor.shape[0]
+    for i in nl.affine_range(math.ceil(num_rows / 128)):
+        row_mask = i * 128 + ix < num_rows
+        a_tile = nl.load(a_tensor[i * 128 + ix, iy], mask=row_mask)
+        max_vals = nl.max(a_tile, axis=[1], keepdims=True, mask=row_mask)
+        shifted = nl.subtract(a_tile, max_vals, mask=row_mask)
+        numerator = nl.exp(shifted)
+        denominator = nl.sum(numerator, axis=[1])  # [128, 1], as in NKI
+        nl.store(
+            out_tensor[i * 128 + ix, iy], value=numerator / denominator, mask=row_mask
+        )
+    return out_tensor
+
+
+def copy_kernel(in_tensor):
+    # Whole-tensor load and store, and a load nested in a store.
+    out_tensor = nl.ndarray(
+        in_tensor.shape, dtype=in_tensor.dtype, buffer=nl.shared_hbm
+    )
+    tmp = nl.ndarray(in_tensor.shape, dtype=in_tensor.dtype, buffer=nl.shared_hbm)
+    nl.store(tmp, value=nl.load(in_tensor))
+    ix = nl.arange(in_tensor.shape[0])[:, None]
+    iy = nl.arange(in_tensor.shape[1])[None, :]
+    nl.store(out_tensor[ix, iy], value=nl.copy(nl.load(tmp[ix, iy])))
+    return out_tensor
+
+
+def spmd_add_kernel(a_input, b_input):
+    # One [128, 256] tile per program on a 2-D grid.
+    c_output = nl.ndarray(a_input.shape, dtype=a_input.dtype, buffer=nl.shared_hbm)
+    ix = nl.program_id(0) * 128 + nl.arange(128)[:, None]
+    iy = nl.program_id(1) * 256 + nl.arange(256)[None, :]
+    nl.store(
+        c_output[ix, iy],
+        value=nl.add(nl.load(a_input[ix, iy]), nl.load(b_input[ix, iy])),
+    )
+    return c_output
+
+
+def loops_and_casts_kernel(x):
+    # static_range / sequential_range loops, par_dim in a shape, static_cast,
+    # maximum / minimum / divide, and a reduction over two free axes.
+    out = nl.ndarray((x.shape[0], 1), dtype=nl.float32, buffer=nl.shared_hbm)
+    acc = nl.zeros((nl.par_dim(x.shape[0]), 1), dtype=nl.float32)
+    for k in nl.static_range(2):
+        t = nl.load(x[:, k])
+        clipped = nl.minimum(nl.maximum(t, -1.0), 1.0)
+        acc[...] = nl.add(acc, nl.sum(clipped, axis=[1, 2]))
+    for _ in nl.sequential_range(1):
+        acc[...] = nl.divide(acc, 2.0)
+    nl.store(out, value=nl.static_cast(acc, nl.float32))
+    return out
+
+
+def _tilelens(kernel, grid, *args):
+    return tilelens.trace(client=Tracer(), frontend="nki")(kernel)[grid](*args)
+
+
+def _reference(kernel, grid, *args):
+    jitted = nki.jit(kernel)
+    return nki.simulate_kernel(jitted[grid] if grid != (1,) else jitted, *args)
+
+
+rng = np.random.default_rng(0)
+
+
+@pytest.mark.parametrize(
+    "kernel, grid, args",
+    [
+        (softmax_kernel, (1,), (rng.standard_normal((200, 64)).astype(np.float32),)),
+        (copy_kernel, (1,), (rng.standard_normal((16, 32)).astype(np.float32),)),
+        (
+            spmd_add_kernel,
+            (2, 2),
+            (
+                rng.standard_normal((256, 512)).astype(np.float32),
+                rng.standard_normal((256, 512)).astype(np.float32),
+            ),
+        ),
+        (
+            loops_and_casts_kernel,
+            (1,),
+            (rng.standard_normal((8, 2, 4, 3)).astype(np.float32),),
+        ),
+    ],
+    ids=["softmax", "copy", "spmd_add", "loops_and_casts"],
+)
+def test_kernel_idioms_match_nki_simulator(kernel, grid, args):
+    tilelens.clear()
+    got = _tilelens(kernel, grid, *args)
+    want = _reference(kernel, grid, *args)
+    assert isinstance(got, np.ndarray)
+    np.testing.assert_allclose(got, want, rtol=1e-5, atol=1e-6)
+
+
+class _StopAfterFirstProgram(Tracer):
+    def post_run_callback(self, fn):
+        return False
+
+
+def test_outputs_are_returned_when_a_client_stops_the_grid_early():
+    tilelens.clear()
+    a = rng.standard_normal((256, 512)).astype(np.float32)
+    b = rng.standard_normal((256, 512)).astype(np.float32)
+    traced = tilelens.trace(client=_StopAfterFirstProgram(), frontend="nki")(
+        spmd_add_kernel
+    )
+    got = traced[(2, 2)](a, b)
+    assert isinstance(got, np.ndarray)
+    # Only program (0, 0) ran; its tile holds the sum.
+    np.testing.assert_allclose(got[:128, :256], a[:128, :256] + b[:128, :256])
diff --git a/tests/nki/test_nki_tracer.py b/tests/nki/test_nki_tracer.py
index ec948bab..8dae4951 100644
--- a/tests/nki/test_nki_tracer.py
+++ b/tests/nki/test_nki_tracer.py
@@ -127,7 +127,8 @@ def reduce_sum_kernel(x_ptr, out_ptr):
     assert record.input_shape == (block_m, block_n)
     assert record.index == 1
     assert record.keep_dims is False
-    assert record.output_shape == (block_m,)
+    # As in NKI, reducing the only free axis of a [P, F] tile keeps it: [P, 1].
+    assert record.output_shape == (block_m, 1)
 
 
 def test_tracer_records_dot():
diff --git a/tilelens/clients/tracer/tracer.py b/tilelens/clients/tracer/tracer.py
index 8151ed95..512d1b50 100644
--- a/tilelens/clients/tracer/tracer.py
+++ b/tilelens/clients/tracer/tracer.py
@@ -13,10 +13,20 @@
     Allocate,
 )
 from ...utils.traceback_utils import extract_user_frames
-from tilelens.core.masked_load_store import masked_load
+from tilelens.core.masked_load_store import masked_load, tile_mask, tile_shape
 import numpy as np
 
 
+def _nki_offsets_and_mask(tensor, keys, mask):
+    """Element offsets and mask of an NKI access ``tensor[keys]``. NKI masks broadcast
+    to the tile shape, and an access without a mask touches every lane."""
+    offsets = tensor.get_offsets().data
+    mask_data = tile_mask(offsets, keys, mask.data if mask is not None else None)
+    if mask_data is None:
+        mask_data = np.ones(tile_shape(offsets, keys), dtype=bool)
+    return masked_load(offsets, keys, mask=mask_data), mask_data
+
+
 def _convert_grid_idx(grid_idx) -> tuple[int, int, int] | None:
     if grid_idx is None:
         return grid_idx
@@ -103,12 +113,13 @@ def pre_load_callback(ptr, mask, keys):
                 first_ptr = np.reshape(ptr.data, (-1))[0]
                 tensor = self._get_tensor(first_ptr)
                 offsets = ptr.data - tensor.data_ptr()
+                mask_data = mask.data
             else:
                 keys = _convert_keys_to_numpy(keys)
-                offsets = masked_load(ptr.get_offsets().data, keys, mask=mask.data)
+                offsets, mask_data = _nki_offsets_and_mask(ptr, keys, mask)
                 tensor = ptr
 
-            rec = Load(tensor.data_ptr(), offsets, mask.data)
+            rec = Load(tensor.data_ptr(), offsets, mask_data)
             rec.call_path = extract_user_frames(num_frames=1)
             self.records.append(rec)
 
@@ -124,12 +135,7 @@ def pre_store_callback(ptr, mask, keys):
                 mask_data = mask.data
             else:
                 keys = _convert_keys_to_numpy(keys)
-                if mask is None:
-                    offsets = masked_load(ptr.get_offsets().data, keys)
-                    mask_data = np.ones_like(offsets).astype(bool)
-                else:
-                    mask_data = mask.data
-                    offsets = masked_load(ptr.get_offsets().data, keys, mask=mask_data)
+                offsets, mask_data = _nki_offsets_and_mask(ptr, keys, mask)
                 tensor = ptr
 
             rec = Store(tensor.data_ptr(), offsets, mask_data)
diff --git a/tilelens/core/frontend/nki.py b/tilelens/core/frontend/nki.py
index 6b0820f1..42e51707 100644
--- a/tilelens/core/frontend/nki.py
+++ b/tilelens/core/frontend/nki.py
@@ -19,7 +19,7 @@
 HAS_NKI = False
 nki_builder = None
 try:
-    from tilelens.core.simulation.nki import NDArray, nki_builder  # type: ignore
+    from tilelens.core.simulation.nki import NDArray, full_keys, nki_builder  # type: ignore
 
     HAS_NKI = True
 except ModuleNotFoundError:
@@ -62,15 +62,21 @@ def _nki_reduce_sum_adapter(
     NKI_ADAPTERS = {
         ProgramId: lambda axis, *_args, **_kwargs: AdapterResult(axis),
         Allocate: lambda *_args, **_kwargs: AdapterResult(),
-        Load: lambda src, keys, *, mask=None, **_kwargs: AdapterResult(
+        # `nl.load(t)` and `nl.store(t, value=v)` access the whole tensor.
+        Load: lambda src, keys=None, *, mask=None, **_kwargs: AdapterResult(
             src,
             mask,
-            keys,
+            full_keys(src) if keys is None else keys,
         ),
-        Store: lambda dst, keys, value, *, mask=None, **_kwargs: AdapterResult(
+        Store: lambda dst,
+        keys=None,
+        value=None,
+        *,
+        mask=None,
+        **_kwargs: AdapterResult(
             dst,
             mask,
-            keys,
+            full_keys(dst) if keys is None else keys,
         ),
         Dot: _nki_dot_adapter,
         ReduceSum: _nki_reduce_sum_adapter,
diff --git a/tilelens/core/frontend/nki_transform.py b/tilelens/core/frontend/nki_transform.py
index 374e2816..cc497fb2 100644
--- a/tilelens/core/frontend/nki_transform.py
+++ b/tilelens/core/frontend/nki_transform.py
@@ -8,35 +8,19 @@ class NKITransformer(ast.NodeTransformer):
 
     Transformations:
     - `nl.<func>(x[...], ...)` -> `nl.<func>(x, slice_obj, ...)`
-       where <func> is one of {load, store, load_transpose2d}
+       where <func> is one of {load, store, load_transpose2d}, wherever the call
+       appears (e.g. nested in another call, as in `nl.store(o[i], nl.load(x[i]))`)
     - strips out kernel decorators
     """
 
-    def visit_Expr(self, node: ast.Expr) -> ast.AST | list[ast.AST]:
-        if not isinstance(node.value, ast.Call):
-            return self.generic_visit(node)
-
-        new_call = self._rewrite_nl_slice_call(node.value)
-        if new_call is None:
-            return self.generic_visit(node)
-        return ast.Expr(value=new_call)
-
     def visit_FunctionDef(self, node: ast.FunctionDef) -> ast.AST:
         node.decorator_list = []
         return self.generic_visit(node)
 
-    def visit_Assign(self, node: ast.Assign) -> ast.AST:
-        if not isinstance(node.value, ast.Call):
-            return self.generic_visit(node)
-
-        new_call = self._rewrite_nl_slice_call(node.value)
-        if new_call is None:
-            return self.generic_visit(node)
-        return ast.Assign(
-            targets=node.targets,
-            value=new_call,
-            type_comment=getattr(node, "type_comment", None),
-        )
+    def visit_Call(self, node: ast.Call) -> ast.AST:
+        node = cast(ast.Call, self.generic_visit(node))
+        new_call = self._rewrite_nl_slice_call(node)
+        return node if new_call is None else new_call
 
     def _rewrite_nl_slice_call(self, call_node: ast.Call) -> ast.Call | None:
         if not (
diff --git a/tilelens/core/masked_load_store.py b/tilelens/core/masked_load_store.py
index f13ff19a..b3c9c732 100644
--- a/tilelens/core/masked_load_store.py
+++ b/tilelens/core/masked_load_store.py
@@ -67,6 +67,30 @@ def _get_valid_indices(
     return valid_mask
 
 
+def tile_shape(ndarray: np.ndarray, keys) -> tuple:
+    """Shape of ``ndarray[keys]``, also when some integer-array keys are out of bounds."""
+    if not isinstance(keys, tuple):
+        keys = (keys,)
+    clipped = []
+    for dim, key in enumerate(keys):
+        if (
+            isinstance(key, np.ndarray)
+            and key.dtype.kind in "iu"
+            and dim < ndarray.ndim
+        ):
+            key = np.clip(key, 0, max(ndarray.shape[dim] - 1, 0))
+        clipped.append(key)
+    return ndarray[tuple(clipped)].shape
+
+
+def tile_mask(ndarray: np.ndarray, keys, mask) -> np.ndarray | None:
+    """``mask`` broadcast to the shape of ``ndarray[keys]``, as NKI broadcasts masks;
+    None stays None."""
+    if mask is None:
+        return None
+    return np.broadcast_to(np.asarray(mask, dtype=bool), tile_shape(ndarray, keys))
+
+
 def masked_load(
     ndarray: np.ndarray, keys: tuple, mask: np.ndarray = None
 ) -> np.ndarray:
diff --git a/tilelens/core/simulation/nki.py b/tilelens/core/simulation/nki.py
index cc1701f6..3a485fa2 100644
--- a/tilelens/core/simulation/nki.py
+++ b/tilelens/core/simulation/nki.py
@@ -12,7 +12,29 @@
 import inspect
 import textwrap
 from ..frontend.nki_transform import transform_code
-from ..masked_load_store import masked_load, masked_store
+from ..masked_load_store import masked_load, masked_store, tile_mask, tile_shape
+
+
+def _value(x):
+    """The numpy value of an NDArray operand; other operands (scalars) as is."""
+    return x.data if isinstance(x, NDArray) else x
+
+
+def _axis(axis):
+    """NKI accepts a list of axes; numpy needs an int or a tuple."""
+    return tuple(axis) if isinstance(axis, list) else axis
+
+
+def _extreme(dtype, lowest):
+    """Identity of a masked max (lowest) or min (highest) reduction over `dtype`."""
+    if np.issubdtype(dtype, np.integer):
+        info = np.iinfo(dtype)
+        return info.min if lowest else info.max
+    return -np.inf if lowest else np.inf
+
+
+def _cast(value, dtype):
+    return np.asarray(value) if dtype is None else np.asarray(value).astype(dtype)
 
 
 class NDArray:
@@ -251,31 +273,48 @@ def _convert_keys_to_numpy(self, keys):
         else:
             return keys
 
-    def load(self, src: NDArray, keys, *, mask=None, **kwargs):
+    def load(self, src: NDArray, keys=None, *, mask=None, **kwargs):
         """Load array elements with masking for out-of-bounds errors."""
         # Convert NDArray to numpy array
         ndarray = src.data
         mask_value = getattr(mask, "data", mask) if mask is not None else None
+        dtype = kwargs.pop("dtype", None)
+        keys = full_keys(src) if keys is None else keys
 
         # Convert any NDArrays in keys to numpy arrays
         numpy_keys = self._convert_keys_to_numpy(keys)
 
         # Call the actual masked_load function
+        mask_value = tile_mask(ndarray, numpy_keys, mask_value)
         result = masked_load(ndarray, numpy_keys, mask=mask_value)
 
         # Convert result back to NDArray
-        return NDArray(value=result, name=f"{src.name}_masked_load", **kwargs)
+        return NDArray(
+            value=_cast(result, dtype), name=f"{src.name}_masked_load", **kwargs
+        )
 
-    def store(self, dst: NDArray, keys, value: NDArray, *, mask=None, **kwargs):
+    def store(
+        self, dst: NDArray, keys=None, value: NDArray = None, *, mask=None, **kwargs
+    ):
         """Store array elements with masking for out-of-bounds errors."""
         # Convert NDArrays to numpy arrays
         ndarray = dst.data
-        value_array = value.data
         mask_value = getattr(mask, "data", mask) if mask is not None else None
+        keys = full_keys(dst) if keys is None else keys
 
         # Convert any NDArrays in keys to numpy arrays
         numpy_keys = self._convert_keys_to_numpy(keys)
 
+        # NKI broadcasts the stored value and the mask to the tile shape
+        tile_shape_ = tile_shape(ndarray, numpy_keys)
+        value_array = np.asarray(_value(value))
+        same_size = value_array.size == int(np.prod(tile_shape_))
+        if value_array.shape != tile_shape_ and same_size:
+            # e.g. a [P, 1] reduction stored into a 1-D slice
+            value_array = value_array.reshape(tile_shape_)
+        value_array = np.broadcast_to(value_array, tile_shape_).astype(ndarray.dtype)
+        mask_value = tile_mask(ndarray, numpy_keys, mask_value)
+
         # Call the actual masked_store function
         masked_store(ndarray, numpy_keys, value_array, mask=mask_value)
 
@@ -290,8 +329,43 @@ def load_transpose2d(
         value = loaded.data.astype(dtype) if dtype is not None else loaded.data
         return NDArray(value=value.T, name=f"{src.name}_load_transpose2d", **kwargs)
 
-    def _unary_op(self, x: NDArray, np_func, op_name, **kwargs):
-        return NDArray(value=np_func(x.data), name=f"{x.name}_{op_name}", **kwargs)
+    def _unary_op(
+        self, x: NDArray, np_func, op_name, *, dtype=None, mask=None, **kwargs
+    ):
+        return NDArray(
+            value=_cast(np_func(x.data), dtype), name=f"{x.name}_{op_name}", **kwargs
+        )
+
+    def _binary_op(self, x, y, np_func, op_name, *, dtype=None, mask=None, **kwargs):
+        return NDArray(
+            value=_cast(np_func(_value(x), _value(y)), dtype),
+            name=f"{getattr(x, 'name', 'scalar')}_{op_name}_{getattr(y, 'name', 'scalar')}",
+            **kwargs,
+        )
+
+    def _reduce(
+        self,
+        x,
+        axis,
+        np_func,
+        op_name,
+        *,
+        dtype=None,
+        mask=None,
+        keepdims=False,
+        initial=None,
+        **kwargs,
+    ):
+        reduce_kwargs = {}
+        if mask is not None:
+            reduce_kwargs["where"] = np.asarray(_value(mask), dtype=bool)
+            if initial is not None:
+                reduce_kwargs["initial"] = initial
+        value = np_func(x.data, axis=_axis(axis), keepdims=keepdims, **reduce_kwargs)
+        if not keepdims and np.ndim(value) == 1 and x.data.ndim >= 2:
+            # NKI keeps a free dimension: reducing all free axes of a [P, ...] tile gives [P, 1]
+            value = value[:, None]
+        return NDArray(value=_cast(value, dtype), name=f"{x.name}_{op_name}", **kwargs)
 
     # Elementwise operator implementations
     def exp(self, x: NDArray, **kwargs):
@@ -359,37 +433,58 @@ def matmul(self, x: NDArray, y: NDArray, transpose_x=False, mask=None, **kwargs)
     def copy(self, x: NDArray, **kwargs):
         return self._unary_op(x, np.copy, "copy", **kwargs)
 
-    def sum(self, x: NDArray, *args, mask=None, **kwargs):
-        if mask is not None:
-            kwargs["where"] = mask.data
+    def sum(self, x: NDArray, *args, **kwargs):
+        axis = args[0] if args else kwargs.pop("axis", None)
+        kwargs.setdefault("keepdims", kwargs.pop("keep_dims", False))
+        return self._reduce(x, axis, np.sum, "sum", **kwargs)
+
+    def max(self, x: NDArray, axis, **kwargs):
+        lowest = _extreme(x.data.dtype, lowest=True)
+        return self._reduce(x, axis, np.max, "max", initial=lowest, **kwargs)
+
+    def min(self, x: NDArray, axis, **kwargs):
+        highest = _extreme(x.data.dtype, lowest=False)
+        return self._reduce(x, axis, np.min, "min", initial=highest, **kwargs)
+
+    def mean(self, x: NDArray, axis, **kwargs):
+        return self._reduce(x, axis, np.mean, "mean", **kwargs)
+
+    def add(self, x, y, **kwargs):
+        return self._binary_op(x, y, np.add, "add", **kwargs)
+
+    def subtract(self, x, y, **kwargs):
+        return self._binary_op(x, y, np.subtract, "subtract", **kwargs)
+
+    def divide(self, x, y, **kwargs):
+        return self._binary_op(x, y, np.divide, "divide", **kwargs)
+
+    def maximum(self, x, y, **kwargs):
+        return self._binary_op(x, y, np.maximum, "maximum", **kwargs)
+
+    def minimum(self, x, y, **kwargs):
+        return self._binary_op(x, y, np.minimum, "minimum", **kwargs)
+
+    def static_cast(self, x, dtype):
         return NDArray(
-            value=x.data.sum(*args, **kwargs), name=f"{x.name}_sum", **kwargs
+            value=_cast(_value(x), dtype), name=f"{getattr(x, 'name', 'scalar')}_cast"
         )
 
+    def par_dim(self, value):
+        # Marks a partition dimension for the compiler; the size is all the simulator needs.
+        return value
+
     def square(self, x: NDArray, **kwargs):
         return self._unary_op(x, np.square, "square", **kwargs)
 
     def rsqrt(self, x: NDArray, **kwargs):
         return self._unary_op(x, lambda v: 1 / np.sqrt(v), "rsqrt", **kwargs)
 
-    def multiply(self, x: NDArray, y: NDArray, **kwargs):
-        if isinstance(y, NDArray):
-            return NDArray(
-                value=np.multiply(x.data, y.data),
-                name=f"{x.name}_multiply_{y.name}",
-                **kwargs,
-            )
-        elif np.isscalar(y):
-            return NDArray(
-                value=np.multiply(x.data, y),
-                name=f"{x.name}_multiply_scalar",
-                **kwargs,
-            )
-        else:
-            raise TypeError(f"Unsupported type for multiply: {type(y)}")
+    def multiply(self, x, y, **kwargs):
+        return self._binary_op(x, y, np.multiply, "multiply", **kwargs)
 
-    def range(self, stop):
-        return range(stop)
+    def range(self, *args, **kwargs):
+        # affine_range, sequential_range, and static_range all iterate like range()
+        return range(*args)
 
 
 nki_builder = Builder()
@@ -415,7 +510,9 @@ def _set_attr(obj, name, value):
     # nl.shared_hbm
     # nl.psum
     _set_attr(nl, "affine_range", nki_builder.range)
-    nl.par_dim
+    _set_attr(nl, "sequential_range", nki_builder.range)
+    _set_attr(nl, "static_range", nki_builder.range)
+    _set_attr(nl, "par_dim", nki_builder.par_dim)
     _set_attr(nl, "zeros", nki_builder.zeros)
     _set_attr(nl, "mgrid", NDArray(value=np.mgrid, buffer=nl.sbuf, name="mgrid"))
     _set_attr(nl, "matmul", nki_builder.matmul)
@@ -424,6 +521,18 @@ def _set_attr(obj, name, value):
     _set_attr(nl, "square", nki_builder.square)
     _set_attr(nl, "rsqrt", nki_builder.rsqrt)
     _set_attr(nl, "multiply", nki_builder.multiply)
+    for name in (
+        "max",
+        "min",
+        "mean",
+        "add",
+        "subtract",
+        "divide",
+        "maximum",
+        "minimum",
+        "static_cast",
+    ):
+        _set_attr(nl, name, getattr(nki_builder, name))
 
     # attention-specific
     _set_attr(nl, "load_transpose2d", nki_builder.load_transpose2d)
@@ -454,6 +563,19 @@ def nki_unpatch_lang(scope=None):
         scope.restore()
 
 
+def full_keys(tensor):
+    """Keys of a whole-tensor access such as `nl.load(t)` or `nl.store(t, value=v)`."""
+    return tuple(slice(None) for _ in range(tensor.data.ndim))
+
+
+def _to_numpy(ret):
+    if isinstance(ret, NDArray):
+        return ret.data
+    if isinstance(ret, (tuple, list)):
+        return type(ret)(_to_numpy(r) for r in ret)
+    return ret
+
+
 class NKIInterpretedFunction:
     def __init__(self, fn):
         self.fn = fn
@@ -520,6 +642,7 @@ def run(self, *args, **kwargs):
             ret = arg
             client_manager.arg_callback(name, arg, ret)
 
+        ret = None
         for x in range(grid_dims[0]):
             for y in range(grid_dims[1]):
                 for z in range(grid_dims[2]):
@@ -530,7 +653,11 @@ def run(self, *args, **kwargs):
                         client_manager.grid_idx_callback((x, y, z))
 
                     if not client_manager.pre_run_callback(self.fn):
-                        return
-                    self.fn(*args, **kwargs)
+                        return _to_numpy(ret)
+                    ret = self.fn(*args, **kwargs)
                     if not client_manager.post_run_callback(self.fn):
-                        return
+                        # a client that needs no more programs stops the grid early
+                        return _to_numpy(ret)
+        # Kernels that allocate their outputs in shared HBM return them; every program
+        # returns the same arrays, so return the last program's, as numpy arrays.
+        return _to_numpy(ret)