From e6e0641f6114dba2b7d1b4e33d6c3c70cf6e91cd Mon Sep 17 00:00:00 2001
From: Jokeren <robinho364@gmail.com>
Date: Mon, 28 Jul 2025 10:09:33 -0400
Subject: [PATCH 01/90] Support NKI

---
 .gitignore               |   3 +
 tests/test_nki.py        |  82 +++++++++++++++++++++++
 triton_viz/core/nki.py   | 138 +++++++++++++++++++++++++++++++++++++++
 triton_viz/core/trace.py |  89 +++++++++++++++----------
 4 files changed, 279 insertions(+), 33 deletions(-)
 create mode 100644 tests/test_nki.py
 create mode 100644 triton_viz/core/nki.py

diff --git a/.gitignore b/.gitignore
index 68bc17f9..060367fa 100644
--- a/.gitignore
+++ b/.gitignore
@@ -3,6 +3,9 @@ __pycache__/
 *.py[cod]
 *$py.class
 
+# VSCode
+.vscode
+
 # C extensions
 *.so
 
diff --git a/tests/test_nki.py b/tests/test_nki.py
new file mode 100644
index 00000000..bd43ecab
--- /dev/null
+++ b/tests/test_nki.py
@@ -0,0 +1,82 @@
+#!/usr/bin/env python3
+"""
+Test script to verify NDArray slicing functionality after fixes
+"""
+import sys
+import os
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'triton_viz', 'core'))
+
+import numpy as np
+
+# Import the nki module directly
+exec(open('triton_viz/core/nki.py').read())
+
+def test_ndarray_creation():
+    print("Testing NDArray creation...")
+    
+    # Test creation with value
+    data = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
+    nd1 = NDArray(value=data, name='test_array')
+    print(f"Created NDArray: {nd1}")
+    print(f"Shape: {nd1.shape}")
+    print(f"Dtype: {nd1.dtype}")
+    print(f"Value:\n{nd1.value}")
+    print()
+    
+    # Test creation with shape and dtype
+    nd2 = NDArray(shape=(2, 3), dtype=np.float32, name='shaped_array')
+    print(f"Created shaped NDArray: {nd2}")
+    print(f"Shape: {nd2.shape}")
+    print(f"Dtype: {nd2.dtype}")
+    print()
+
+def test_slicing():
+    print("Testing slicing operations...")
+    
+    data = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
+    nd_array = NDArray(value=data, name='test_array')
+    
+    # Test [:, :] (all elements)
+    slice_all = nd_array[:, :]
+    print(f"nd_array[:, :] = {slice_all}")
+    print(f"Value:\n{slice_all.value}")
+    print()
+    
+    # Test [:, 0] (first column)
+    slice_col = nd_array[:, 0]
+    print(f"nd_array[:, 0] = {slice_col}")
+    print(f"Value: {slice_col.value}")
+    print()
+    
+    # Test [0, :] (first row)
+    slice_row = nd_array[0, :]
+    print(f"nd_array[0, :] = {slice_row}")
+    print(f"Value: {slice_row.value}")
+    print()
+
+def test_arithmetic():
+    print("Testing arithmetic operations...")
+    
+    data1 = np.array([[1, 2], [3, 4]])
+    data2 = np.array([[5, 6], [7, 8]])
+    
+    nd1 = NDArray(value=data1, name='array1')
+    nd2 = NDArray(value=data2, name='array2')
+    
+    # Test addition
+    result = nd1 + nd2
+    print(f"Addition result: {result}")
+    print(f"Value:\n{result.value}")
+    print()
+    
+    # Test slicing on result
+    slice_result = result[:, 0]
+    print(f"Slice of result [:, 0]: {slice_result}")
+    print(f"Value: {slice_result.value}")
+    print()
+
+if __name__ == "__main__":
+    test_ndarray_creation()
+    test_slicing()
+    test_arithmetic()
+    print("All tests completed successfully!")
diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
new file mode 100644
index 00000000..400469dc
--- /dev/null
+++ b/triton_viz/core/nki.py
@@ -0,0 +1,138 @@
+import numpy as np
+
+
+class NDArray:
+    def __init__(self, buffer=None, name="", **kwargs):
+        self.buffer = buffer
+        self.name = name
+        self.kwargs = kwargs
+        if "shape" in kwargs and "dtype" in kwargs:
+            shape = kwargs.pop("shape")
+            dtype = kwargs.pop("dtype")
+            self._value = np.ndarray(shape, dtype=dtype, buffer=buffer, **kwargs)
+        elif "value" in kwargs:
+            self._value = kwargs["value"]
+        else:
+            self._value = None
+
+    @property
+    def shape(self):
+        return self._value.shape if self._value is not None else None
+
+    @property
+    def dtype(self):
+        return self._value.dtype if self._value is not None else None
+
+    @property
+    def value(self):
+        return self._value
+
+    @value.setter
+    def value(self, new_value):
+        self._value = new_value
+
+    def __repr__(self):
+        return f"NDArray(shape={self.shape}, dtype={self.dtype}, name={self.name})"
+
+    def __getitem__(self, key):
+        """Implement slicing operations for NDArray"""
+        if self._value is None:
+            raise AttributeError("NDArray has no value to slice")
+
+        # Apply the slicing to the underlying numpy array
+        sliced_value = self._value[key]
+
+        # Create a new NDArray with the sliced data
+        return NDArray(value=sliced_value, name=f"{self.name}_slice")
+
+    # Define operator +/-/*//
+    def __add__(self, other):
+        if isinstance(other, NDArray):
+            return NDArray(value=self._value + other._value, name=f"{self.name}_add_{other.name}")
+        raise TypeError(f"Unsupported operand type(s) for +: 'NDArray' and '{type(other).__name__}'")
+
+    def __sub__(self, other):
+        if isinstance(other, NDArray):
+            return NDArray(value=self._value - other._value, name=f"{self.name}_sub_{other.name}")
+        raise TypeError(f"Unsupported operand type(s) for -: 'NDArray' and '{type(other).__name__}'")
+
+    def __mul__(self, other):
+        if isinstance(other, NDArray):
+            return NDArray(value=self._value * other._value, name=f"{self.name}_mul_{other.name}")
+        raise TypeError(f"Unsupported operand type(s) for *: 'NDArray' and '{type(other).__name__}'")
+
+    def __truediv__(self, other):
+        if isinstance(other, NDArray):
+            return NDArray(value=self._value / other._value, name=f"{self.name}_div_{other.name}")
+        raise TypeError(f"Unsupported operand type(s) for /: 'NDArray' and '{type(other).__name__}'")
+
+
+class Builder:
+    def __init__(self, grid_dims=None):
+        self.grid_dims = grid_dims if grid_dims is not None else (1, 1, 1)
+        self.grid_x = None
+        self.grid_y = None
+        self.grid_z = None
+
+    def set_grid_dims(self, grid_dims):
+        self.grid_dims = grid_dims
+
+    def set_grid_idx(self, x, y, z):
+        self.grid_x = x
+        self.grid_y = y
+        self.grid_z = z
+
+    def ndarray(self, shape, dtype, *, buffer=None, name="", **kwargs):
+        return NDArray(buffer=buffer, name=name, shape=shape, dtype=dtype, **kwargs)
+
+    def arange(self, *args):
+        if len(args) == 1:
+            value = np.arange(args[0])
+        elif len(args) == 2:
+            value = np.arange(args[0], args[1])
+        else:
+            raise ValueError("arange expects 1 or 2 arguments")
+        return NDArray(value=value, name="arange")
+
+    def program_id(self, axis: int):
+        if axis == 0:
+            return self.grid_x
+        elif axis == 1:
+            return self.grid_y
+        elif axis == 2:
+            return self.grid_z
+        else:
+            raise ValueError(f"Invalid axis: {axis}. Must be 0, 1, or 2.")
+
+    def load(self, src: NDArray, *, mask=None, dtype=None, **kwargs):
+        value = src._value
+        if mask is not None:
+            value = value[mask]
+        if dtype is not None:
+            value = value.astype(dtype)
+        return NDArray(value=value, name=src.name, **kwargs)
+
+    def store(self, dst: NDArray, value: NDArray, *, mask=None, **kwargs):
+        if mask is not None:
+            value = value[mask]
+        dst._value = value._value
+        return dst
+
+
+nki_builder = Builder()
+
+
+class NKIInterpreterFunction:
+    def __init__(self, fn):
+        self.fn = fn
+
+    def run(self, *args, **kwargs):
+        grid_dims = kwargs.pop("grid", (1, 1, 1))  # Remove grid from kwargs to avoid passing it to the function
+        nki_builder.set_grid_dims(grid_dims)
+        
+        for x in range(grid_dims[0]):
+            for y in range(grid_dims[1]):
+                for z in range(grid_dims[2]):
+                    nki_builder.set_grid_idx(x, y, z)
+                    result = self.fn(*args, **kwargs)
+        return result
diff --git a/triton_viz/core/trace.py b/triton_viz/core/trace.py
index fa002d07..f26f6829 100644
--- a/triton_viz/core/trace.py
+++ b/triton_viz/core/trace.py
@@ -6,12 +6,20 @@
 from ..clients import Sanitizer, Profiler, Tracer
 from .client import ClientManager, Client
 from .data import Launch
+from typing import TypeVar, Generic
 
 
 launches: list[Launch] = []
 
+T = TypeVar("T")
+
+
+class TraceInterface(Generic[T]):
+    def __init__(self, client: str | Client, interpreter_fn: T) -> None:
+        self.client_manager = ClientManager()
+        self.add_client(client)
+        self.interpreter_fn: T = interpreter_fn
 
-class Trace(KernelInterface):
     @staticmethod
     def _normalize_client(client: str | Client) -> Client:
         if isinstance(client, str):
@@ -31,40 +39,54 @@ def _normalize_client(client: str | Client) -> Client:
     def add_client(self, new_client: str | Client) -> None:
         self.client_manager.add_clients([self._normalize_client(new_client)])
 
+    def finalize(self):
+        self.client_manager.finalize()
+        launches.append(self.client_manager.launch)
+
+    def run(self, *args, **kwargs):
+        with self.client_manager.patch():
+            kwargs.update({"client_manager": self.client_manager})
+            ret = self.interpreter_fn.run(*args, **kwargs)
+            self.finalize()
+            return ret
+
+
+class TritonTrace(KernelInterface, TraceInterface):
     def __init__(
         self,
         kernel: JITFunction | InterpretedFunction,
         client: str | Client,
     ) -> None:
-        self.fn = kernel
         if isinstance(kernel, InterpretedFunction):
-            self.interpreter_fn = kernel
+            interpreter_fn = kernel
         elif isinstance(kernel, JITFunction):
-            self.interpreter_fn = InterpretedFunction(kernel.fn)
+            interpreter_fn = InterpretedFunction(kernel.fn)
         else:
-            raise TypeError(
-                f"Kernel must be JITFunction or InterpretedFunction, got {type(kernel)}"
-            )
+            raise TypeError(f"Kernel must be JITFunction or InterpretedFunction, got {type(kernel)}")
+        TraceInterface.__init__(self, client, interpreter_fn)
+        self.fn = (kernel,)
         self.arg_names = kernel.arg_names
-        self.client_manager = ClientManager()
-        self.add_client(client)
-
-    def run(self, *args, **kwargs):
-        with self.client_manager.patch():
-            kwargs.update({"client_manager": self.client_manager})
-            ret = self.interpreter_fn.run(*args, **kwargs)
-            self.finalize()
-            return ret
 
     def warmup(self, *args, **kwargs):
         raise NotImplementedError
 
-    def finalize(self):
-        self.client_manager.finalize()
-        launches.append(self.client_manager.launch)
+
+class NKITrace(KernelInterface, TraceInterface):
+    def __init__(self, kernel, client: str | Client) -> None:
+        from neuronxcc.nki.compile import GenericKernel
+        from .nki import NKIInterpretedFunction
+
+        if isinstance(kernel, GenericKernel):
+            interpreter_fn = NKIInterpretedFunction(kernel.func)
+        elif isinstance(kernel, NKIInterpretedFunction):
+            interpreter_fn = kernel
+        TraceInterface.__init__(self, client, interpreter_fn)
+
+    def __getitem__(self, *grid):
+        return KernelInterface.__getitem__(self, tuple(*grid))
 
 
-def trace(clients: str | Client | None = None):
+def trace(clients: str | Client | None = None, backend: str = "triton"):
     """
     Create a trace object that can be used to run a kernel with instrumentation clients.
 
@@ -77,24 +99,25 @@ def trace(clients: str | Client | None = None):
     if not isinstance(clients, (str, Client)):
         raise TypeError(f"Expected str or Client, got {type(clients)}")
 
-    def decorator(kernel) -> Trace:
+    def decorator(kernel) -> TritonTrace:
         # When sanitizer is off, skip tracing and return the original kernel unchanged
         if cfg.sanitizer_backend == "off":
             return kernel
 
         # First-time wrapping
-        if isinstance(kernel, (JITFunction, InterpretedFunction)):
-            return Trace(kernel, clients)
-
-        # If the object is already a Trace, just append the new client(s)
-        if isinstance(kernel, Trace):
-            trace = kernel
-            trace.add_client(clients)
-            return trace
-
-        raise TypeError(
-            f"Expected JITFunction, InterpretedFunction or Trace, got {type(kernel)}"
-        )
+        if backend == "triton":
+            if isinstance(kernel, (JITFunction, InterpretedFunction)):
+                return TritonTrace(kernel, clients)
+
+            # If the object is already a TritonTrace, just append the new client(s)
+            if isinstance(kernel, TritonTrace):
+                trace = kernel
+                trace.add_client(clients)
+                return trace
+        elif backend == "nki":
+            return NKITrace(kernel, clients)
+
+        raise TypeError(f"Expected JITFunction, InterpretedFunction or Trace, got {type(kernel)}")
 
     return decorator
 

From c9b2501dc8cce2a2013e68af46780c3055f3ea51 Mon Sep 17 00:00:00 2001
From: Jokeren <robinho364@gmail.com>
Date: Mon, 28 Jul 2025 10:16:01 -0400
Subject: [PATCH 02/90] Update

---
 triton_viz/core/nki.py   | 2 +-
 triton_viz/core/trace.py | 3 ++-
 2 files changed, 3 insertions(+), 2 deletions(-)

diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index 400469dc..f5a82875 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -122,7 +122,7 @@ def store(self, dst: NDArray, value: NDArray, *, mask=None, **kwargs):
 nki_builder = Builder()
 
 
-class NKIInterpreterFunction:
+class NKIInterpretedFunction:
     def __init__(self, fn):
         self.fn = fn
 
diff --git a/triton_viz/core/trace.py b/triton_viz/core/trace.py
index f26f6829..69f2bb03 100644
--- a/triton_viz/core/trace.py
+++ b/triton_viz/core/trace.py
@@ -64,7 +64,7 @@ def __init__(
         else:
             raise TypeError(f"Kernel must be JITFunction or InterpretedFunction, got {type(kernel)}")
         TraceInterface.__init__(self, client, interpreter_fn)
-        self.fn = (kernel,)
+        self.fn = kernel
         self.arg_names = kernel.arg_names
 
     def warmup(self, *args, **kwargs):
@@ -80,6 +80,7 @@ def __init__(self, kernel, client: str | Client) -> None:
             interpreter_fn = NKIInterpretedFunction(kernel.func)
         elif isinstance(kernel, NKIInterpretedFunction):
             interpreter_fn = kernel
+        self.func = kernel.func
         TraceInterface.__init__(self, client, interpreter_fn)
 
     def __getitem__(self, *grid):

From 6fe57e6f6a33e86019edf57418c57a93fe9a4f5c Mon Sep 17 00:00:00 2001
From: Jokeren <robinho364@gmail.com>
Date: Mon, 28 Jul 2025 10:18:00 -0400
Subject: [PATCH 03/90] Update

---
 triton_viz/core/nki.py | 7 +++++++
 1 file changed, 7 insertions(+)

diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index f5a82875..5ebe1603 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -128,6 +128,13 @@ def __init__(self, fn):
 
     def run(self, *args, **kwargs):
         grid_dims = kwargs.pop("grid", (1, 1, 1))  # Remove grid from kwargs to avoid passing it to the function
+        # make it 3d if not
+        if len(grid_dims) == 1:
+            grid_dims = (grid_dims[0], 1, 1)
+        elif len(grid_dims) == 2:
+            grid_dims = (grid_dims[0], grid_dims[1], 1)
+        elif len(grid_dims) != 3:
+            raise ValueError(f"Grid must be 1, 2, or 3 dimensions, got {len(grid_dims)}")
         nki_builder.set_grid_dims(grid_dims)
         
         for x in range(grid_dims[0]):

From 86d1d6f1d5a902a30706e146bbf3cfde45a630e5 Mon Sep 17 00:00:00 2001
From: Jokeren <robinho364@gmail.com>
Date: Mon, 28 Jul 2025 10:18:41 -0400
Subject: [PATCH 04/90] Update

---
 triton_viz/core/nki.py | 2 ++
 1 file changed, 2 insertions(+)

diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index 5ebe1603..55f803e6 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -136,6 +136,8 @@ def run(self, *args, **kwargs):
         elif len(grid_dims) != 3:
             raise ValueError(f"Grid must be 1, 2, or 3 dimensions, got {len(grid_dims)}")
         nki_builder.set_grid_dims(grid_dims)
+
+        kwargs.pop("warmup", None)  # Remove warmup from kwargs if it exists
         
         for x in range(grid_dims[0]):
             for y in range(grid_dims[1]):

From 7aa3d9130a35308255f885bb3ac366b2c4061c79 Mon Sep 17 00:00:00 2001
From: Jokeren <robinho364@gmail.com>
Date: Mon, 28 Jul 2025 10:19:24 -0400
Subject: [PATCH 05/90] Update

---
 triton_viz/core/nki.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index 55f803e6..77b2efd9 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -138,6 +138,7 @@ def run(self, *args, **kwargs):
         nki_builder.set_grid_dims(grid_dims)
 
         kwargs.pop("warmup", None)  # Remove warmup from kwargs if it exists
+        kwargs.pop("client_manager", None)  # Remove client_manager from kwargs if it exists
         
         for x in range(grid_dims[0]):
             for y in range(grid_dims[1]):

From 4f8f616300decd1c2fe026a515896e0268c77dcd Mon Sep 17 00:00:00 2001
From: Jokeren <robinho364@gmail.com>
Date: Mon, 28 Jul 2025 10:30:47 -0400
Subject: [PATCH 06/90] Update

---
 triton_viz/core/nki.py | 21 ++++++++++++++++++++-
 1 file changed, 20 insertions(+), 1 deletion(-)

diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index 77b2efd9..6790855c 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -1,5 +1,7 @@
 import numpy as np
 
+import neuronxcc.nki.language as nl
+
 
 class NDArray:
     def __init__(self, buffer=None, name="", **kwargs):
@@ -122,6 +124,21 @@ def store(self, dst: NDArray, value: NDArray, *, mask=None, **kwargs):
 nki_builder = Builder()
 
 
+def patch():
+    nl.ndarray = lambda *args, **kwargs: nki_builder.ndarray(*args, **kwargs)
+    nl.program_id = lambda axis: nki_builder.program_id(axis)
+    nl.arange = lambda *args: nki_builder.arange(*args)
+    nl.load = lambda src, **kwargs: nki_builder.load(src, **kwargs)
+    nl.store = lambda dst, value, **kwargs: nki_builder.store(dst, value, **kwargs)
+
+
+def unpatch():
+    # reload the original functions
+    import importlib
+
+    importlib.reload(nl)
+
+
 class NKIInterpretedFunction:
     def __init__(self, fn):
         self.fn = fn
@@ -139,10 +156,12 @@ def run(self, *args, **kwargs):
 
         kwargs.pop("warmup", None)  # Remove warmup from kwargs if it exists
         kwargs.pop("client_manager", None)  # Remove client_manager from kwargs if it exists
-        
+
+        patch()
         for x in range(grid_dims[0]):
             for y in range(grid_dims[1]):
                 for z in range(grid_dims[2]):
                     nki_builder.set_grid_idx(x, y, z)
                     result = self.fn(*args, **kwargs)
+        unpatch()
         return result

From 9962f53cca1b31a9bc57b0da4aa7f9ed4c6c2c1b Mon Sep 17 00:00:00 2001
From: Jokeren <robinho364@gmail.com>
Date: Mon, 28 Jul 2025 10:33:10 -0400
Subject: [PATCH 07/90] Update

---
 triton_viz/core/trace.py | 6 +++++-
 1 file changed, 5 insertions(+), 1 deletion(-)

diff --git a/triton_viz/core/trace.py b/triton_viz/core/trace.py
index 69f2bb03..49e23e2d 100644
--- a/triton_viz/core/trace.py
+++ b/triton_viz/core/trace.py
@@ -77,10 +77,14 @@ def __init__(self, kernel, client: str | Client) -> None:
         from .nki import NKIInterpretedFunction
 
         if isinstance(kernel, GenericKernel):
+            # This is wrong
             interpreter_fn = NKIInterpretedFunction(kernel.func)
+            self.func = kernel.func
         elif isinstance(kernel, NKIInterpretedFunction):
             interpreter_fn = kernel
-        self.func = kernel.func
+        else:
+            interpreter_fn = NKIInterpretedFunction(kernel.func)
+            self.func = kernel
         TraceInterface.__init__(self, client, interpreter_fn)
 
     def __getitem__(self, *grid):

From 25c8cebfb58e8bda5e8adec4e0390ffc46c16f2c Mon Sep 17 00:00:00 2001
From: Jokeren <robinho364@gmail.com>
Date: Mon, 28 Jul 2025 10:33:35 -0400
Subject: [PATCH 08/90] Update

---
 triton_viz/core/trace.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/triton_viz/core/trace.py b/triton_viz/core/trace.py
index 49e23e2d..8d75d658 100644
--- a/triton_viz/core/trace.py
+++ b/triton_viz/core/trace.py
@@ -83,7 +83,7 @@ def __init__(self, kernel, client: str | Client) -> None:
         elif isinstance(kernel, NKIInterpretedFunction):
             interpreter_fn = kernel
         else:
-            interpreter_fn = NKIInterpretedFunction(kernel.func)
+            interpreter_fn = NKIInterpretedFunction(kernel)
             self.func = kernel
         TraceInterface.__init__(self, client, interpreter_fn)
 

From cb0783112d8ac8198893c31880d269eb0bddee0b Mon Sep 17 00:00:00 2001
From: Jokeren <robinho364@gmail.com>
Date: Mon, 28 Jul 2025 10:34:39 -0400
Subject: [PATCH 09/90] Update

---
 triton_viz/core/nki.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index 6790855c..95006b43 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -11,7 +11,7 @@ def __init__(self, buffer=None, name="", **kwargs):
         if "shape" in kwargs and "dtype" in kwargs:
             shape = kwargs.pop("shape")
             dtype = kwargs.pop("dtype")
-            self._value = np.ndarray(shape, dtype=dtype, buffer=buffer, **kwargs)
+            self._value = np.ndarray(shape, dtype=dtype)
         elif "value" in kwargs:
             self._value = kwargs["value"]
         else:

From 181028497659e3fffa4a4c11ca16a64c6c595528 Mon Sep 17 00:00:00 2001
From: Jokeren <robinho364@gmail.com>
Date: Mon, 28 Jul 2025 10:36:10 -0400
Subject: [PATCH 10/90] Update

---
 triton_viz/core/nki.py | 8 ++++++++
 1 file changed, 8 insertions(+)

diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index 95006b43..442c1cfa 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -51,21 +51,29 @@ def __getitem__(self, key):
     def __add__(self, other):
         if isinstance(other, NDArray):
             return NDArray(value=self._value + other._value, name=f"{self.name}_add_{other.name}")
+        elif np.isscalar(other):
+            return NDArray(value=self._value + other, name=f"{self.name}_add_scalar")
         raise TypeError(f"Unsupported operand type(s) for +: 'NDArray' and '{type(other).__name__}'")
 
     def __sub__(self, other):
         if isinstance(other, NDArray):
             return NDArray(value=self._value - other._value, name=f"{self.name}_sub_{other.name}")
+        elif np.isscalar(other):
+            return NDArray(value=self._value - other, name=f"{self.name}_sub_scalar")
         raise TypeError(f"Unsupported operand type(s) for -: 'NDArray' and '{type(other).__name__}'")
 
     def __mul__(self, other):
         if isinstance(other, NDArray):
             return NDArray(value=self._value * other._value, name=f"{self.name}_mul_{other.name}")
+        elif np.isscalar(other):
+            return NDArray(value=self._value * other, name=f"{self.name}_mul_scalar")
         raise TypeError(f"Unsupported operand type(s) for *: 'NDArray' and '{type(other).__name__}'")
 
     def __truediv__(self, other):
         if isinstance(other, NDArray):
             return NDArray(value=self._value / other._value, name=f"{self.name}_div_{other.name}")
+        elif np.isscalar(other):
+            return NDArray(value=self._value / other, name=f"{self.name}_div_scalar")
         raise TypeError(f"Unsupported operand type(s) for /: 'NDArray' and '{type(other).__name__}'")
 
 

From e6a6f919be59bf9ac8cf3109a52c600fa293ec94 Mon Sep 17 00:00:00 2001
From: Jokeren <robinho364@gmail.com>
Date: Mon, 28 Jul 2025 10:36:52 -0400
Subject: [PATCH 11/90] Update

---
 triton_viz/core/nki.py | 28 ++++++++++++++++++++++++++++
 1 file changed, 28 insertions(+)

diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index 442c1cfa..6a49967a 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -55,6 +55,13 @@ def __add__(self, other):
             return NDArray(value=self._value + other, name=f"{self.name}_add_scalar")
         raise TypeError(f"Unsupported operand type(s) for +: 'NDArray' and '{type(other).__name__}'")
 
+    def __radd__(self, other):
+        if isinstance(other, NDArray):
+            return NDArray(value=other._value + self._value, name=f"{other.name}_add_{self.name}")
+        elif np.isscalar(other):
+            return NDArray(value=other + self._value, name=f"scalar_add_{self.name}")
+        raise TypeError(f"Unsupported operand type(s) for +: '{type(other).__name__}' and 'NDArray'")
+
     def __sub__(self, other):
         if isinstance(other, NDArray):
             return NDArray(value=self._value - other._value, name=f"{self.name}_sub_{other.name}")
@@ -62,6 +69,13 @@ def __sub__(self, other):
             return NDArray(value=self._value - other, name=f"{self.name}_sub_scalar")
         raise TypeError(f"Unsupported operand type(s) for -: 'NDArray' and '{type(other).__name__}'")
 
+    def __rsub__(self, other):
+        if isinstance(other, NDArray):
+            return NDArray(value=other._value - self._value, name=f"{other.name}_sub_{self.name}")
+        elif np.isscalar(other):
+            return NDArray(value=other - self._value, name=f"scalar_sub_{self.name}")
+        raise TypeError(f"Unsupported operand type(s) for -: '{type(other).__name__}' and 'NDArray'")
+
     def __mul__(self, other):
         if isinstance(other, NDArray):
             return NDArray(value=self._value * other._value, name=f"{self.name}_mul_{other.name}")
@@ -69,6 +83,13 @@ def __mul__(self, other):
             return NDArray(value=self._value * other, name=f"{self.name}_mul_scalar")
         raise TypeError(f"Unsupported operand type(s) for *: 'NDArray' and '{type(other).__name__}'")
 
+    def __rmul__(self, other):
+        if isinstance(other, NDArray):
+            return NDArray(value=other._value * self._value, name=f"{other.name}_mul_{self.name}")
+        elif np.isscalar(other):
+            return NDArray(value=other * self._value, name=f"scalar_mul_{self.name}")
+        raise TypeError(f"Unsupported operand type(s) for *: '{type(other).__name__}' and 'NDArray'")
+
     def __truediv__(self, other):
         if isinstance(other, NDArray):
             return NDArray(value=self._value / other._value, name=f"{self.name}_div_{other.name}")
@@ -76,6 +97,13 @@ def __truediv__(self, other):
             return NDArray(value=self._value / other, name=f"{self.name}_div_scalar")
         raise TypeError(f"Unsupported operand type(s) for /: 'NDArray' and '{type(other).__name__}'")
 
+    def __rtruediv__(self, other):
+        if isinstance(other, NDArray):
+            return NDArray(value=other._value / self._value, name=f"{other.name}_div_{self.name}")
+        elif np.isscalar(other):
+            return NDArray(value=other / self._value, name=f"scalar_div_{self.name}")
+        raise TypeError(f"Unsupported operand type(s) for /: '{type(other).__name__}' and 'NDArray'")
+
 
 class Builder:
     def __init__(self, grid_dims=None):

From 8406f667889e457d072b3ea198a92a4dfe22faa6 Mon Sep 17 00:00:00 2001
From: Jokeren <robinho364@gmail.com>
Date: Mon, 28 Jul 2025 10:39:57 -0400
Subject: [PATCH 12/90] Update

---
 triton_viz/core/nki.py | 11 +++++++++--
 1 file changed, 9 insertions(+), 2 deletions(-)

diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index 6a49967a..bf5e9f92 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -36,13 +36,20 @@ def value(self, new_value):
     def __repr__(self):
         return f"NDArray(shape={self.shape}, dtype={self.dtype}, name={self.name})"
 
-    def __getitem__(self, key):
+    def __getitem__(self, keys):
         """Implement slicing operations for NDArray"""
         if self._value is None:
             raise AttributeError("NDArray has no value to slice")
 
         # Apply the slicing to the underlying numpy array
-        sliced_value = self._value[key]
+        new_keys = []
+        if isinstance(keys, tuple):
+            for k in keys:
+                if isinstance(k, NDArray):
+                    new_keys.append(k._value)
+                else:
+                    new_keys.append(k)
+        sliced_value = self._value[tuple(new_keys)]
 
         # Create a new NDArray with the sliced data
         return NDArray(value=sliced_value, name=f"{self.name}_slice")

From b63e74c178f578895e9614c1ac796070cae3ffa5 Mon Sep 17 00:00:00 2001
From: Jokeren <robinho364@gmail.com>
Date: Mon, 28 Jul 2025 10:44:40 -0400
Subject: [PATCH 13/90] Update

---
 triton_viz/core/nki.py | 4 ++++
 1 file changed, 4 insertions(+)

diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index bf5e9f92..775ee3a6 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -54,6 +54,7 @@ def __getitem__(self, keys):
         # Create a new NDArray with the sliced data
         return NDArray(value=sliced_value, name=f"{self.name}_slice")
 
+
     # Define operator +/-/*//
     def __add__(self, other):
         if isinstance(other, NDArray):
@@ -201,6 +202,9 @@ def run(self, *args, **kwargs):
         kwargs.pop("client_manager", None)  # Remove client_manager from kwargs if it exists
 
         patch()
+
+        # convert args to NDArray if they are not already
+        args = [arg if isinstance(arg, NDArray) else NDArray(value=arg) for arg in args]
         for x in range(grid_dims[0]):
             for y in range(grid_dims[1]):
                 for z in range(grid_dims[2]):

From 241530bc3f29c4e1a5dcf1c7f143282af1d2bdbf Mon Sep 17 00:00:00 2001
From: Jokeren <robinho364@gmail.com>
Date: Mon, 28 Jul 2025 10:45:40 -0400
Subject: [PATCH 14/90] Update

---
 triton_viz/core/nki.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index 775ee3a6..9c02b71a 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -211,4 +211,4 @@ def run(self, *args, **kwargs):
                     nki_builder.set_grid_idx(x, y, z)
                     result = self.fn(*args, **kwargs)
         unpatch()
-        return result
+        return result.value

From d28a17eb29196ff7893004ac385c879f52de00f9 Mon Sep 17 00:00:00 2001
From: Jokeren <robinho364@gmail.com>
Date: Mon, 28 Jul 2025 10:50:05 -0400
Subject: [PATCH 15/90] Update

---
 triton_viz/core/nki.py | 14 +++++++++++++-
 1 file changed, 13 insertions(+), 1 deletion(-)

diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index 9c02b71a..eb60c0ae 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -1,6 +1,7 @@
 import numpy as np
 
 import neuronxcc.nki.language as nl
+import inspect
 
 
 class NDArray:
@@ -114,11 +115,13 @@ def __rtruediv__(self, other):
 
 
 class Builder:
-    def __init__(self, grid_dims=None):
+    def __init__(self, fn, grid_dims=None):
+        self.fn = fn
         self.grid_dims = grid_dims if grid_dims is not None else (1, 1, 1)
         self.grid_x = None
         self.grid_y = None
         self.grid_z = None
+        self.shared_hbm_arrays = {}
 
     def set_grid_dims(self, grid_dims):
         self.grid_dims = grid_dims
@@ -129,6 +132,15 @@ def set_grid_idx(self, x, y, z):
         self.grid_z = z
 
     def ndarray(self, shape, dtype, *, buffer=None, name="", **kwargs):
+        if name is None:
+            # file name + function name + line number
+            frame = inspect.currentframe().f_back
+            file_name = frame.f_code.co_filename
+            function_name = frame.f_code.co_name
+            line_number = frame.f_lineno
+            name = f"{file_name}_{function_name}_{line_number}"
+        if buffer == nl.shared_hbm and name in self.shared_hbm_arrays:
+            return self.shared_hbm_arrays[name]
         return NDArray(buffer=buffer, name=name, shape=shape, dtype=dtype, **kwargs)
 
     def arange(self, *args):

From bc0873af49e5a707378076acc013d79032b8906d Mon Sep 17 00:00:00 2001
From: Jokeren <robinho364@gmail.com>
Date: Mon, 28 Jul 2025 10:51:38 -0400
Subject: [PATCH 16/90] Update

---
 triton_viz/core/nki.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index eb60c0ae..c2c22d88 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -115,12 +115,12 @@ def __rtruediv__(self, other):
 
 
 class Builder:
-    def __init__(self, fn, grid_dims=None):
-        self.fn = fn
+    def __init__(self, grid_dims=None):
         self.grid_dims = grid_dims if grid_dims is not None else (1, 1, 1)
         self.grid_x = None
         self.grid_y = None
         self.grid_z = None
+        self.fn = None
         self.shared_hbm_arrays = {}
 
     def set_grid_dims(self, grid_dims):
@@ -209,6 +209,8 @@ def run(self, *args, **kwargs):
         elif len(grid_dims) != 3:
             raise ValueError(f"Grid must be 1, 2, or 3 dimensions, got {len(grid_dims)}")
         nki_builder.set_grid_dims(grid_dims)
+        nki_builder.shared_hbm_arrays = {}
+        nki_builder.fn = self.fn
 
         kwargs.pop("warmup", None)  # Remove warmup from kwargs if it exists
         kwargs.pop("client_manager", None)  # Remove client_manager from kwargs if it exists

From 4b6827d51688c1d21335d3ac1d52178acf666d75 Mon Sep 17 00:00:00 2001
From: Jokeren <robinho364@gmail.com>
Date: Mon, 28 Jul 2025 10:54:02 -0400
Subject: [PATCH 17/90] UPdate

---
 triton_viz/core/nki.py | 14 +++++++++++---
 1 file changed, 11 insertions(+), 3 deletions(-)

diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index c2c22d88..9a0d7190 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -139,9 +139,17 @@ def ndarray(self, shape, dtype, *, buffer=None, name="", **kwargs):
             function_name = frame.f_code.co_name
             line_number = frame.f_lineno
             name = f"{file_name}_{function_name}_{line_number}"
-        if buffer == nl.shared_hbm and name in self.shared_hbm_arrays:
-            return self.shared_hbm_arrays[name]
-        return NDArray(buffer=buffer, name=name, shape=shape, dtype=dtype, **kwargs)
+        if buffer == nl.shared_hbm:
+            if name in self.shared_hbm_arrays:
+                # Return the existing shared HBM array
+                ret = self.shared_hbm_arrays[name]
+            else:
+                # Create a new shared HBM array and store it
+                ret = NDArray(buffer=buffer, name=name, shape=shape, dtype=dtype, **kwargs)
+                self.shared_hbm_arrays[name] = ret
+        else:
+            ret = NDArray(buffer=buffer, name=name, shape=shape, dtype=dtype, **kwargs)
+        return ret
 
     def arange(self, *args):
         if len(args) == 1:

From 62ade2dfab356e2974e8dc4e79e1701139d8b1d6 Mon Sep 17 00:00:00 2001
From: Jokeren <robinho364@gmail.com>
Date: Mon, 28 Jul 2025 10:54:47 -0400
Subject: [PATCH 18/90] Update

---
 triton_viz/core/nki.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index 9a0d7190..61677b3e 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -132,14 +132,14 @@ def set_grid_idx(self, x, y, z):
         self.grid_z = z
 
     def ndarray(self, shape, dtype, *, buffer=None, name="", **kwargs):
-        if name is None:
-            # file name + function name + line number
-            frame = inspect.currentframe().f_back
-            file_name = frame.f_code.co_filename
-            function_name = frame.f_code.co_name
-            line_number = frame.f_lineno
-            name = f"{file_name}_{function_name}_{line_number}"
         if buffer == nl.shared_hbm:
+            if name is None:
+                # file name + function name + line number
+                frame = inspect.currentframe().f_back
+                file_name = frame.f_code.co_filename
+                function_name = frame.f_code.co_name
+                line_number = frame.f_lineno
+                name = f"{file_name}_{function_name}_{line_number}"
             if name in self.shared_hbm_arrays:
                 # Return the existing shared HBM array
                 ret = self.shared_hbm_arrays[name]

From d9d7bf7779582d97ba7ea0e5b8f2e2a68953a792 Mon Sep 17 00:00:00 2001
From: Jokeren <robinho364@gmail.com>
Date: Mon, 28 Jul 2025 11:53:08 -0400
Subject: [PATCH 19/90] Update

---
 triton_viz/core/nki.py | 55 +++++++++++++++++++++++++++++++++++++++---
 1 file changed, 52 insertions(+), 3 deletions(-)

diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index 61677b3e..e4b9720b 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -4,6 +4,48 @@
 import inspect
 
 
+class NLSlice:
+    def __init__(self, start=None, stop=None, step: int = 1):
+        self.start = start
+        self.stop = stop
+        self.step = step
+
+    def __repr__(self):
+        return f"NLSlice(start={self.start}, stop={self.stop}, step={self.step})"
+
+    def to_tuple(self):
+        return (self.start, self.stop, self.step) if self.step is not None else (self.start, self.stop)
+
+    def __add__(self, other):
+        if isinstance(other, NLSlice):
+            return NLSlice(
+                start=self.start + other.start if self.start is not None else None,
+                stop=self.stop + other.stop if self.stop is not None else None,
+                step=self.step + other.step if self.step is not None else None
+            )
+        elif isinstance(other, int):
+            return NLSlice(
+                start=self.start + other if self.start is not None else None,
+                stop=self.stop + other if self.stop is not None else None,
+                step=self.step
+            )
+        raise TypeError(f"Unsupported operand type(s) for +: 'NLSlice' and '{type(other).__name__}'")
+
+    def __radd__(self, other):
+        if isinstance(other, NLSlice):
+            return NLSlice(
+                start=other.start + self.start if self.start is not None else None,
+                stop=other.stop + self.stop if self.stop is not None else None,
+                step=other.step + self.step if self.step is not None else None
+            )
+        elif isinstance(other, int):
+            return NLSlice(
+                start=other + self.start if self.start is not None else None,
+                stop=other + self.stop if self.stop is not None else None,
+                step=self.step
+            )
+        raise TypeError(f"Unsupported operand type(s) for +: '{type(other).__name__}' and 'NLSlice'")
+
 class NDArray:
     def __init__(self, buffer=None, name="", **kwargs):
         self.buffer = buffer
@@ -48,6 +90,8 @@ def __getitem__(self, keys):
             for k in keys:
                 if isinstance(k, NDArray):
                     new_keys.append(k._value)
+                elif isinstance(k, NLSlice):
+                    new_keys.append(slice(k.start, k.stop, k.step))
                 else:
                     new_keys.append(k)
         sliced_value = self._value[tuple(new_keys)]
@@ -153,12 +197,17 @@ def ndarray(self, shape, dtype, *, buffer=None, name="", **kwargs):
 
     def arange(self, *args):
         if len(args) == 1:
-            value = np.arange(args[0])
+            start = 0
+            stop = args[0]
         elif len(args) == 2:
-            value = np.arange(args[0], args[1])
+            start = args[0]
+            stop = args[1]
         else:
             raise ValueError("arange expects 1 or 2 arguments")
-        return NDArray(value=value, name="arange")
+        return NLSlice(
+            start=start,
+            stop=stop,
+        ) 
 
     def program_id(self, axis: int):
         if axis == 0:

From 998a1a51d79f075f06b4dc17fd96c46b05b0db33 Mon Sep 17 00:00:00 2001
From: Jokeren <robinho364@gmail.com>
Date: Mon, 28 Jul 2025 12:18:38 -0400
Subject: [PATCH 20/90] Update

---
 triton_viz/core/nki.py | 86 +++++++++++++++++++++++++++++-------------
 1 file changed, 59 insertions(+), 27 deletions(-)

diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index e4b9720b..a3e3c37c 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -4,48 +4,80 @@
 import inspect
 
 
+# Multi-dimensional slice class
 class NLSlice:
     def __init__(self, start=None, stop=None, step: int = 1):
-        self.start = start
-        self.stop = stop
-        self.step = step
+        self.start = [start] if isinstance(start, int) else start
+        self.stop = [stop] if isinstance(stop, int) else stop
+        self.step = [step] if isinstance(step, int) else step
 
     def __repr__(self):
-        return f"NLSlice(start={self.start}, stop={self.stop}, step={self.step})"
-
-    def to_tuple(self):
-        return (self.start, self.stop, self.step) if self.step is not None else (self.start, self.stop)
+        repr = ""
+        for start, stop, step in zip(self.start, self.stop, self.step):
+            if start is None:
+                start = "None"
+            if stop is None:
+                stop = "None"
+            if step is None:
+                step = "None"
+            repr += f"(start={start}, stop={stop}, step={step}) " 
+        return repr
 
     def __add__(self, other):
+        new_start = []
+        new_stop = []
+        new_step = []
         if isinstance(other, NLSlice):
-            return NLSlice(
-                start=self.start + other.start if self.start is not None else None,
-                stop=self.stop + other.stop if self.stop is not None else None,
-                step=self.step + other.step if self.step is not None else None
-            )
+            for start, stop, step in zip(self.start, self.stop, self.step):
+                new_start.append(start + other.start if start is not None else None)
+                new_stop.append(stop + other.stop if stop is not None else None)
+                new_step.append(step + other.step if step is not None else None)
+            return NLSlice(start=new_start, stop=new_stop, step=new_step)
         elif isinstance(other, int):
-            return NLSlice(
-                start=self.start + other if self.start is not None else None,
-                stop=self.stop + other if self.stop is not None else None,
-                step=self.step
-            )
+            for start, stop, step in zip(self.start, self.stop, self.step):
+                new_start.append(start + other if start is not None else None)
+                new_stop.append(stop + other if stop is not None else None)
+                new_step.append(step)
+            return NLSlice(start=new_start, stop=new_stop, step=new_step)
         raise TypeError(f"Unsupported operand type(s) for +: 'NLSlice' and '{type(other).__name__}'")
 
     def __radd__(self, other):
+        new_start = []
+        new_stop = []
+        new_step = []
         if isinstance(other, NLSlice):
-            return NLSlice(
-                start=other.start + self.start if self.start is not None else None,
-                stop=other.stop + self.stop if self.stop is not None else None,
-                step=other.step + self.step if self.step is not None else None
-            )
+            for start, stop, step in zip(self.start, self.stop, self.step):
+                new_start.append(other.start + start if start is not None else None)
+                new_stop.append(other.stop + stop if stop is not None else None)
+                new_step.append(other.step + step if step is not None else None)
+            return NLSlice(start=new_start, stop=new_stop, step=new_step)
         elif isinstance(other, int):
-            return NLSlice(
-                start=other + self.start if self.start is not None else None,
-                stop=other + self.stop if self.stop is not None else None,
-                step=self.step
-            )
+            for start, stop, step in zip(self.start, self.stop, self.step):
+                new_start.append(other + start if start is not None else None)
+                new_stop.append(other + stop if stop is not None else None)
+                new_step.append(step)
+            return NLSlice(start=new_start, stop=new_stop, step=new_step)
         raise TypeError(f"Unsupported operand type(s) for +: '{type(other).__name__}' and 'NLSlice'")
 
+    def __getitem__(self, keys):
+        new_start = []
+        new_stop = []
+        new_step = []
+        for k in keys:
+            # check if k is None:
+            if k is None:
+                new_start.append(None)
+                new_stop.append(None)
+                new_step.append(None)
+            elif isinstance(k, slice):
+                new_start.append(k.start)
+                new_stop.append(k.stop)
+                new_step.append(k.step)
+            else:
+                raise TypeError(f"Unsupported key type: {type(k)}")
+
+        return NLSlice(start=new_start, stop=new_stop, step=new_step)
+
 class NDArray:
     def __init__(self, buffer=None, name="", **kwargs):
         self.buffer = buffer

From 2a50a044f991bada0cc559a1357f1f62cc1cdb6c Mon Sep 17 00:00:00 2001
From: Jokeren <robinho364@gmail.com>
Date: Mon, 28 Jul 2025 12:25:36 -0400
Subject: [PATCH 21/90] Update

---
 triton_viz/core/nki.py | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index a3e3c37c..49d19c98 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -63,6 +63,7 @@ def __getitem__(self, keys):
         new_start = []
         new_stop = []
         new_step = []
+        idx = 0
         for k in keys:
             # check if k is None:
             if k is None:
@@ -70,9 +71,11 @@ def __getitem__(self, keys):
                 new_stop.append(None)
                 new_step.append(None)
             elif isinstance(k, slice):
-                new_start.append(k.start)
-                new_stop.append(k.stop)
-                new_step.append(k.step)
+                assert k.start is None and k.stop is None and k.step is None, "Slice must be complete"
+                new_start.append(self.start[idx])
+                new_stop.append(self.stop[idx])
+                new_step.append(self.step[idx])
+                idx += 1
             else:
                 raise TypeError(f"Unsupported key type: {type(k)}")
 

From 8e49f9d8ffbf3d5ff908b20e2303bf1ecb51d07a Mon Sep 17 00:00:00 2001
From: Jokeren <robinho364@gmail.com>
Date: Mon, 28 Jul 2025 12:29:33 -0400
Subject: [PATCH 22/90] Update

---
 triton_viz/core/nki.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index 49d19c98..55d5be0b 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -35,7 +35,7 @@ def __add__(self, other):
             return NLSlice(start=new_start, stop=new_stop, step=new_step)
         elif isinstance(other, int):
             for start, stop, step in zip(self.start, self.stop, self.step):
-                new_start.append(start + other if start is not None else None)
+                new_start.append(start + other if start is not None else other)
                 new_stop.append(stop + other if stop is not None else None)
                 new_step.append(step)
             return NLSlice(start=new_start, stop=new_stop, step=new_step)
@@ -53,7 +53,7 @@ def __radd__(self, other):
             return NLSlice(start=new_start, stop=new_stop, step=new_step)
         elif isinstance(other, int):
             for start, stop, step in zip(self.start, self.stop, self.step):
-                new_start.append(other + start if start is not None else None)
+                new_start.append(other + start if start is not None else other)
                 new_stop.append(other + stop if stop is not None else None)
                 new_step.append(step)
             return NLSlice(start=new_start, stop=new_stop, step=new_step)

From 45596d3c55f53d6e74ca5e162a21c7a2cd8ac2eb Mon Sep 17 00:00:00 2001
From: Jokeren <robinho364@gmail.com>
Date: Mon, 28 Jul 2025 12:31:21 -0400
Subject: [PATCH 23/90] Update

---
 triton_viz/core/nki.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index 55d5be0b..49d19c98 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -35,7 +35,7 @@ def __add__(self, other):
             return NLSlice(start=new_start, stop=new_stop, step=new_step)
         elif isinstance(other, int):
             for start, stop, step in zip(self.start, self.stop, self.step):
-                new_start.append(start + other if start is not None else other)
+                new_start.append(start + other if start is not None else None)
                 new_stop.append(stop + other if stop is not None else None)
                 new_step.append(step)
             return NLSlice(start=new_start, stop=new_stop, step=new_step)
@@ -53,7 +53,7 @@ def __radd__(self, other):
             return NLSlice(start=new_start, stop=new_stop, step=new_step)
         elif isinstance(other, int):
             for start, stop, step in zip(self.start, self.stop, self.step):
-                new_start.append(other + start if start is not None else other)
+                new_start.append(other + start if start is not None else None)
                 new_stop.append(other + stop if stop is not None else None)
                 new_step.append(step)
             return NLSlice(start=new_start, stop=new_stop, step=new_step)

From a54ac94b0f61244dcef85b5366ec35367f7c24e7 Mon Sep 17 00:00:00 2001
From: Jokeren <robinho364@gmail.com>
Date: Mon, 28 Jul 2025 12:42:25 -0400
Subject: [PATCH 24/90] Update

---
 triton_viz/core/nki.py | 19 +++++++++++++++++++
 1 file changed, 19 insertions(+)

diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index 49d19c98..c48e46dc 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -129,8 +129,27 @@ def __getitem__(self, keys):
                     new_keys.append(slice(k.start, k.stop, k.step))
                 else:
                     new_keys.append(k)
+        # new_keys=[slice([0, None], [128, None], [1, None]), slice([None, 0], [None, 512], [None, 1])]
+        if all(isinstance(k, slice) for k in new_keys):
+
+            def coalesce(*args):
+                for arg in args:
+                    if arg is not None:
+                        return arg
+                return None
+
+            # Combine slices per axis, preferring non-None values
+            new_keys = tuple(
+                slice(
+                    coalesce(*(k.start[i] for k in new_keys)),
+                    coalesce(*(k.stop[i] for k in new_keys)),
+                    coalesce(*(k.step[i] for k in new_keys)),
+                )
+                for i in range(len(new_keys[0].start))
+            )
         sliced_value = self._value[tuple(new_keys)]
 
+
         # Create a new NDArray with the sliced data
         return NDArray(value=sliced_value, name=f"{self.name}_slice")
 

From f1231feb265d5427cb88340da672445ad381f8ca Mon Sep 17 00:00:00 2001
From: Jokeren <robinho364@gmail.com>
Date: Mon, 28 Jul 2025 12:46:02 -0400
Subject: [PATCH 25/90] Update

---
 triton_viz/core/nki.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index c48e46dc..40ebc469 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -284,7 +284,7 @@ def load(self, src: NDArray, *, mask=None, dtype=None, **kwargs):
     def store(self, dst: NDArray, value: NDArray, *, mask=None, **kwargs):
         if mask is not None:
             value = value[mask]
-        dst._value = value._value
+        dst._value[:] = value._value[:]
         return dst
 
 

From b1f658ab90deb52a851b862e30d0400cfd556651 Mon Sep 17 00:00:00 2001
From: Jokeren <robinho364@gmail.com>
Date: Mon, 28 Jul 2025 12:47:41 -0400
Subject: [PATCH 26/90] Update comments

---
 triton_viz/core/nki.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index 40ebc469..30f6e408 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -4,6 +4,9 @@
 import inspect
 
 
+# Q1: slicing semantic is weird
+# Q2: why cannot we execute the .func function?
+
 # Multi-dimensional slice class
 class NLSlice:
     def __init__(self, start=None, stop=None, step: int = 1):

From e7e2dd25c76409fde6e869b2661156eaa9834379 Mon Sep 17 00:00:00 2001
From: latentCall145 <latentCall145>
Date: Wed, 10 Sep 2025 14:30:59 +0000
Subject: [PATCH 27/90] stuff

---
 examples/nki_softmax.py  | 63 ++++++++++++++++++++++++++++++++++++++++
 tests/test_nki.py        |  3 +-
 triton_viz/core/nki.py   | 34 +++++++++++-----------
 triton_viz/core/trace.py |  3 +-
 4 files changed, 84 insertions(+), 19 deletions(-)
 create mode 100644 examples/nki_softmax.py

diff --git a/examples/nki_softmax.py b/examples/nki_softmax.py
new file mode 100644
index 00000000..fc6d9e3c
--- /dev/null
+++ b/examples/nki_softmax.py
@@ -0,0 +1,63 @@
+from neuronxcc import nki
+import neuronxcc.nki.language as nl
+
+import torch
+import triton_viz
+import torch_xla.core.xla_model as xm
+from triton_viz.clients import Tracer
+from triton_viz.core import config as cfg
+from triton_viz.core.trace import launches
+
+
+@triton_viz.trace(clients=Tracer())
+@nki.jit
+def softmax(in_tensor):
+    # assume B,D; softmax on D for now
+    B, D = in_tensor.shape
+    out_tensor = nl.ndarray((B, D), dtype=in_tensor.dtype, buffer=nl.shared_hbm)
+
+    #assert nl.tile_size.pmax == 128
+    num_tiles = math.ceil(B / 128)
+    for tile_idx in nl.affine_range(num_tiles):
+        i_p = tile_idx * 128 + nl.arange(128)[:, None]
+        i_f = nl.arange(D)[None, :]
+        mask = i_p < B
+        tile = nl.load(in_tensor[i_p, i_f], mask=mask)
+        tile_exp = nl.exp(tile)
+        tile_expsum = nl.sum(tile_exp, -1, keepdims=True, mask=mask)
+        tile_softmax = tile_exp / tile_expsum
+        nl.store(out_tensor[i_p, i_f], value=tile_softmax, mask=mask)
+    return out_tensor
+
+
+if __name__ == "__main__":
+    cfg.reset()
+    device = "xla"
+    size = 16
+    BLOCK_SIZE = 8
+    torch.manual_seed(0)
+    x = torch.arange(size, dtype=torch.float32, device=device)
+    softmax[(1,1,1)](x)
+
+    # Print records to see what's being captured
+    print(f"Number of launches: {len(launches)}")
+    if launches:
+        launch = launches[-1]
+        print(f"Number of records: {len(launch.records)}")
+        for i, record in enumerate(launch.records):
+            print(f"Record {i}: {type(record).__name__}")
+            if hasattr(record, "ptr"):
+                print(f"  ptr: {record.ptr}")
+            if hasattr(record, "offsets"):
+                print(f"  offsets shape: {record.offsets.shape}")
+            if hasattr(record, "masks"):
+                print(f"  masks shape: {record.masks.shape}")
+
+    # Try to launch visualization
+    try:
+        triton_viz.launch()
+    except Exception as e:
+        print(f"\nError during visualization: {e}")
+        import traceback
+
+        traceback.print_exc()
diff --git a/tests/test_nki.py b/tests/test_nki.py
index bd43ecab..221d690d 100644
--- a/tests/test_nki.py
+++ b/tests/test_nki.py
@@ -37,7 +37,8 @@ def test_slicing():
     nd_array = NDArray(value=data, name='test_array')
     
     # Test [:, :] (all elements)
-    slice_all = nd_array[:, :]
+    #slice_all = nd_array[:, :]
+    slice_all = nd_array[:2, :2]
     print(f"nd_array[:, :] = {slice_all}")
     print(f"Value:\n{slice_all.value}")
     print()
diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index 30f6e408..bea06e46 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -133,23 +133,23 @@ def __getitem__(self, keys):
                 else:
                     new_keys.append(k)
         # new_keys=[slice([0, None], [128, None], [1, None]), slice([None, 0], [None, 512], [None, 1])]
-        if all(isinstance(k, slice) for k in new_keys):
-
-            def coalesce(*args):
-                for arg in args:
-                    if arg is not None:
-                        return arg
-                return None
-
-            # Combine slices per axis, preferring non-None values
-            new_keys = tuple(
-                slice(
-                    coalesce(*(k.start[i] for k in new_keys)),
-                    coalesce(*(k.stop[i] for k in new_keys)),
-                    coalesce(*(k.step[i] for k in new_keys)),
-                )
-                for i in range(len(new_keys[0].start))
-            )
+        #if all(isinstance(k, slice) for k in new_keys):
+
+        #    def coalesce(*args):
+        #        for arg in args:
+        #            if arg is not None:
+        #                return arg
+        #        return None
+
+        #    # Combine slices per axis, preferring non-None values
+        #    new_keys = tuple(
+        #        slice(
+        #            coalesce(*(k.start[i] for k in new_keys)),
+        #            coalesce(*(k.stop[i] for k in new_keys)),
+        #            coalesce(*(k.step[i] for k in new_keys)),
+        #        )
+        #        for i in range(len(new_keys[0].start))
+        #    )
         sliced_value = self._value[tuple(new_keys)]
 
 
diff --git a/triton_viz/core/trace.py b/triton_viz/core/trace.py
index 8d75d658..e842a134 100644
--- a/triton_viz/core/trace.py
+++ b/triton_viz/core/trace.py
@@ -91,7 +91,8 @@ def __getitem__(self, *grid):
         return KernelInterface.__getitem__(self, tuple(*grid))
 
 
-def trace(clients: str | Client | None = None, backend: str = "triton"):
+#def trace(clients: str | Client | None = None, backend: str = "triton"):
+def trace(clients: str | Client | None = None, backend: str = "nki"):
     """
     Create a trace object that can be used to run a kernel with instrumentation clients.
 

From b3ee8d9cb2625a2f82fe5a87e414346a3553f62a Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Wed, 10 Sep 2025 13:35:52 -0400
Subject: [PATCH 28/90] triton-viz should support cpu to allow development on
 non-trainium node

---
 examples/nki_softmax.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/nki_softmax.py b/examples/nki_softmax.py
index fc6d9e3c..f84b2711 100644
--- a/examples/nki_softmax.py
+++ b/examples/nki_softmax.py
@@ -32,7 +32,7 @@ def softmax(in_tensor):
 
 if __name__ == "__main__":
     cfg.reset()
-    device = "xla"
+    device = "cpu"
     size = 16
     BLOCK_SIZE = 8
     torch.manual_seed(0)

From 454ac737d155819f3388eea2387761192d6d8d1d Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Fri, 26 Sep 2025 04:24:27 -0400
Subject: [PATCH 29/90] test kernels to check basic tracing

---
 nki-examples/softmax.py      |  63 +++++++++++++++
 nki-examples/test_kernels.py | 150 +++++++++++++++++++++++++++++++++++
 2 files changed, 213 insertions(+)
 create mode 100644 nki-examples/softmax.py
 create mode 100644 nki-examples/test_kernels.py

diff --git a/nki-examples/softmax.py b/nki-examples/softmax.py
new file mode 100644
index 00000000..f84b2711
--- /dev/null
+++ b/nki-examples/softmax.py
@@ -0,0 +1,63 @@
+from neuronxcc import nki
+import neuronxcc.nki.language as nl
+
+import torch
+import triton_viz
+import torch_xla.core.xla_model as xm
+from triton_viz.clients import Tracer
+from triton_viz.core import config as cfg
+from triton_viz.core.trace import launches
+
+
+@triton_viz.trace(clients=Tracer())
+@nki.jit
+def softmax(in_tensor):
+    # assume B,D; softmax on D for now
+    B, D = in_tensor.shape
+    out_tensor = nl.ndarray((B, D), dtype=in_tensor.dtype, buffer=nl.shared_hbm)
+
+    #assert nl.tile_size.pmax == 128
+    num_tiles = math.ceil(B / 128)
+    for tile_idx in nl.affine_range(num_tiles):
+        i_p = tile_idx * 128 + nl.arange(128)[:, None]
+        i_f = nl.arange(D)[None, :]
+        mask = i_p < B
+        tile = nl.load(in_tensor[i_p, i_f], mask=mask)
+        tile_exp = nl.exp(tile)
+        tile_expsum = nl.sum(tile_exp, -1, keepdims=True, mask=mask)
+        tile_softmax = tile_exp / tile_expsum
+        nl.store(out_tensor[i_p, i_f], value=tile_softmax, mask=mask)
+    return out_tensor
+
+
+if __name__ == "__main__":
+    cfg.reset()
+    device = "cpu"
+    size = 16
+    BLOCK_SIZE = 8
+    torch.manual_seed(0)
+    x = torch.arange(size, dtype=torch.float32, device=device)
+    softmax[(1,1,1)](x)
+
+    # Print records to see what's being captured
+    print(f"Number of launches: {len(launches)}")
+    if launches:
+        launch = launches[-1]
+        print(f"Number of records: {len(launch.records)}")
+        for i, record in enumerate(launch.records):
+            print(f"Record {i}: {type(record).__name__}")
+            if hasattr(record, "ptr"):
+                print(f"  ptr: {record.ptr}")
+            if hasattr(record, "offsets"):
+                print(f"  offsets shape: {record.offsets.shape}")
+            if hasattr(record, "masks"):
+                print(f"  masks shape: {record.masks.shape}")
+
+    # Try to launch visualization
+    try:
+        triton_viz.launch()
+    except Exception as e:
+        print(f"\nError during visualization: {e}")
+        import traceback
+
+        traceback.print_exc()
diff --git a/nki-examples/test_kernels.py b/nki-examples/test_kernels.py
new file mode 100644
index 00000000..d4426918
--- /dev/null
+++ b/nki-examples/test_kernels.py
@@ -0,0 +1,150 @@
+import math
+import neuronxcc.nki.language as nl
+import neuronxcc.nki as nki
+import numpy as np
+import torch
+from triton_viz.clients import Tracer
+import triton_viz
+
+TRITON_VIZ = True
+
+def add_kernel(a, b): # fails @ ix < B
+    pid_x = nl.program_id(0)
+    pid_y = nl.program_id(1)
+    B, D = a.shape
+
+    # memory shared across all SPMD instances
+    c_output = nl.ndarray(a.shape, buffer=nl.shared_hbm, dtype=a.dtype)
+
+    ix = pid_x * 128 + nl.arange(128)[:, None]
+    iy = pid_y * 512 + nl.arange(512)[None, :]
+    mask = (ix < B) & (iy < D)
+    a_tmp = nl.load(a[ix, iy], mask=mask)
+    b_tmp = nl.load(b[ix, iy], mask=mask)
+    c_tmp = a_tmp + b_tmp
+    nl.store(c_output[ix, iy], value=c_tmp, mask=mask)
+    return c_output
+
+def copy_kernel(a): # fails @ ix < B
+    B, D = a.shape
+    out = nl.ndarray(a.shape, dtype=a.dtype, buffer=nl.shared_hbm)
+
+    pid_x = nl.program_id(0)
+    pid_y = nl.program_id(1)
+
+    ix = pid_x * 128 + nl.arange(128)[:, None]
+    iy = pid_y * 512 + nl.arange(512)[None, :]
+    #mask = (ix < B) & (iy < D)
+
+    #a_tmp = nl.load(a[ix, iy], mask=mask)
+    #nl.store(out[ix, iy], value=a_tmp, mask=mask)
+    a_tmp = nl.load(a[ix, iy])
+    nl.store(out[ix, iy], value=a_tmp)
+    return out
+
+def print_kernel(): # works
+    a = nl.ndarray([4, 4], dtype=nl.float32, buffer=nl.shared_hbm)
+    y = nl.ndarray([4, 4], dtype=np.float32,)
+    nl.store(a, value=y) 
+    print(a)
+    return a
+
+def tmp0_kernel(a): # works
+    B, D = a.shape
+    out = nl.ndarray(a.shape, dtype=a.dtype, buffer=nl.shared_hbm)
+
+    pid_x = nl.program_id(0)
+    pid_y = nl.program_id(1)
+    import tp
+    tp.log(f'{pid_x=}, {pid_y=}')
+
+    ix = pid_x * 128 + nl.arange(128)[:, None]
+    iy = pid_y * 512 + nl.arange(512)[None, :]
+    #mask = (ix < B) & (iy < D)
+
+    #a_tmp = nl.load(a[ix, iy], mask=mask)
+    #nl.store(out[ix, iy], value=a_tmp, mask=mask)
+    #a_tmp = nl.load(a[:128, :128])
+    a_tmp = nl.load(a[ix, iy])
+    nl.store(out[ix, iy], value=a_tmp)
+    return out
+
+def xyz_kernel(a): # works
+    B, T, C, H, W = a.shape
+    out = nl.ndarray(a.shape, dtype=a.dtype, buffer=nl.shared_hbm)
+
+    pid_x = nl.program_id(0)
+    pid_y = nl.program_id(1)
+    pid_z = nl.program_id(2)
+
+
+    for i in range(H // 128):
+        for j in range(W // 512):
+            #i_h = pid_x * 128 + nl.arange(128)[None, None, None, :, None]
+            #i_n = pid_y * 512 + nl.arange(512)[None, None, None, None, :]
+            i_h = pid_x * 128 + nl.arange(128)[:, None]
+            i_n = pid_y * 512 + nl.arange(512)[None, :]
+            a_tmp = nl.load(a[pid_x, pid_y, pid_z, i_h, i_n])
+            nl.store(out[pid_x, pid_y, pid_z, i_h, i_n], value=a_tmp)
+    return out
+
+def xyz_kernel(a): # works
+    B, T, C, H, W = a.shape
+    out = nl.ndarray(a.shape, dtype=a.dtype, buffer=nl.shared_hbm)
+
+    pid_x = nl.program_id(0)
+    pid_y = nl.program_id(1)
+    pid_z = nl.program_id(2)
+
+
+    for i in range(H // 128):
+        for j in range(W // 512):
+            #i_h = pid_x * 128 + nl.arange(128)[None, None, None, :, None]
+            #i_n = pid_y * 512 + nl.arange(512)[None, None, None, None, :]
+            i_h = pid_x * 128 + nl.arange(128)[:, None]
+            i_n = pid_y * 512 + nl.arange(512)[None, :]
+            a_tmp = nl.load(a[pid_x, pid_y, pid_z, i_h, i_n])
+            nl.store(out[pid_x, pid_y, pid_z, i_h, i_n], value=a_tmp)
+    return out
+
+B, D = 1024, 1024
+x = torch.rand((B, D))
+y = torch.rand((B, D))
+
+blocks_x = math.ceil(B / 128)
+blocks_y = math.ceil(D / 512)
+
+kernel = tmp0_kernel
+if kernel == add_kernel:
+    kernel_grid = (blocks_x, blocks_y)
+    kernel_args = (x.numpy(), y.numpy())
+    z1 = x + y
+if kernel == copy_kernel:
+    kernel_grid = (blocks_x, blocks_y)
+    kernel_args = (x.numpy(),)
+    z1 = x
+elif kernel == print_kernel:
+    kernel_grid = (1,1,1)
+    kernel_args = ()
+    z1 = x
+if kernel == tmp0_kernel:
+    kernel_grid = (blocks_x, blocks_y)
+    kernel_args = (x.numpy(),)
+    z1 = x
+if kernel == xyz_kernel:
+    B, T, C, H, W = 2, 3, 4, 1024, 1024
+    x = torch.rand((B, T, C, H, W))
+    kernel_grid = (B, T, C)
+    kernel_args = (x.numpy(),)
+    z1 = x
+
+if TRITON_VIZ:
+    kernel = triton_viz.trace(clients=Tracer(), backend='nki')(kernel)
+    z2 = kernel[kernel_grid](*kernel_args)
+    z2 = torch.from_numpy(z2)
+    print((z1 - z2).abs().max())
+else:
+    kernel = nki.jit(kernel)
+    z2 = nki.simulate_kernel(kernel[kernel_grid], *kernel_args)
+    z2 = torch.from_numpy(z2)
+    print((z1 - z2).abs().max())

From 65e861d49605f29ed3e3df0f186d2eb0f2c36871 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Fri, 26 Sep 2025 04:26:40 -0400
Subject: [PATCH 30/90] more flexible advanced indexing (just make arange
 slices an arange)

---
 tests/test_nki.py      | 33 ++++++++++++++++++++++
 triton_viz/core/nki.py | 62 ++++++++++++++++++++++++++----------------
 2 files changed, 72 insertions(+), 23 deletions(-)

diff --git a/tests/test_nki.py b/tests/test_nki.py
index 221d690d..3f45c47c 100644
--- a/tests/test_nki.py
+++ b/tests/test_nki.py
@@ -39,22 +39,55 @@ def test_slicing():
     # Test [:, :] (all elements)
     #slice_all = nd_array[:, :]
     slice_all = nd_array[:2, :2]
+    assert np.allclose(slice_all._value, data[:2, :2])
     print(f"nd_array[:, :] = {slice_all}")
     print(f"Value:\n{slice_all.value}")
     print()
     
     # Test [:, 0] (first column)
     slice_col = nd_array[:, 0]
+    assert np.allclose(slice_col._value, data[:, 0])
     print(f"nd_array[:, 0] = {slice_col}")
     print(f"Value: {slice_col.value}")
     print()
     
     # Test [0, :] (first row)
     slice_row = nd_array[0, :]
+    assert np.allclose(slice_row._value, data[0, :])
     print(f"nd_array[0, :] = {slice_row}")
     print(f"Value: {slice_row.value}")
     print()
 
+    # Test advanced indexing
+    rewritten_slice = (
+        np.arange(2)[:, None],
+        np.arange(3)[None, :],
+    )
+    slice_advanced = nd_array[rewritten_slice]
+    assert np.allclose(slice_advanced._value, data[rewritten_slice])
+    print(f"nd_array[nl.arange(2)[:, None], nl.arange(3)[None, :]] = {slice_advanced}")
+    print(f"Value:\n{slice_advanced.value}")
+    print()
+
+    # test nd_array[0, :3, :, 2:4, 2]
+    data = np.reshape(np.arange(3 * 4 * 5 * 6 * 7), (3, 4, 5, 6, 7))
+    nd_array = NDArray(value=data, name='test_array_5d')
+    rewritten_slice = (
+        0,
+        np.arange(3)[:, None, None],
+        np.arange(5)[None, :, None],
+        np.arange(2, 4)[None, None, :],
+        2
+    ) # this is what the above slice would be represented as after triton-viz tracing
+    slice_advanced = nd_array[rewritten_slice]
+    assert np.allclose(
+        slice_advanced._value,
+        data[rewritten_slice]
+    )
+    assert np.allclose(slice_advanced._value, data[0, :3, :, 2:4, 2])
+    print(f"Value:\n{slice_advanced.value}")
+    print()
+
 def test_arithmetic():
     print("Testing arithmetic operations...")
     
diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index bea06e46..2a2efaa9 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -132,27 +132,8 @@ def __getitem__(self, keys):
                     new_keys.append(slice(k.start, k.stop, k.step))
                 else:
                     new_keys.append(k)
-        # new_keys=[slice([0, None], [128, None], [1, None]), slice([None, 0], [None, 512], [None, 1])]
-        #if all(isinstance(k, slice) for k in new_keys):
-
-        #    def coalesce(*args):
-        #        for arg in args:
-        #            if arg is not None:
-        #                return arg
-        #        return None
-
-        #    # Combine slices per axis, preferring non-None values
-        #    new_keys = tuple(
-        #        slice(
-        #            coalesce(*(k.start[i] for k in new_keys)),
-        #            coalesce(*(k.stop[i] for k in new_keys)),
-        #            coalesce(*(k.step[i] for k in new_keys)),
-        #        )
-        #        for i in range(len(new_keys[0].start))
-        #    )
         sliced_value = self._value[tuple(new_keys)]
 
-
         # Create a new NDArray with the sliced data
         return NDArray(value=sliced_value, name=f"{self.name}_slice")
 
@@ -217,6 +198,7 @@ def __rtruediv__(self, other):
 
 class Builder:
     def __init__(self, grid_dims=None):
+        # TODO: infinite grid dims for NKI
         self.grid_dims = grid_dims if grid_dims is not None else (1, 1, 1)
         self.grid_x = None
         self.grid_y = None
@@ -251,6 +233,12 @@ def ndarray(self, shape, dtype, *, buffer=None, name="", **kwargs):
         else:
             ret = NDArray(buffer=buffer, name=name, shape=shape, dtype=dtype, **kwargs)
         return ret
+    
+    def zeros(self, shape, dtype, *, buffer=None, name="", **kwargs):
+        val = np.zeros(shape, dtype=dtype)
+        return self.ndarray(
+            shape, dtype, buffer=buffer, name=name, value=value, **kwargs
+        )
 
     def arange(self, *args):
         if len(args) == 1:
@@ -261,10 +249,7 @@ def arange(self, *args):
             stop = args[1]
         else:
             raise ValueError("arange expects 1 or 2 arguments")
-        return NLSlice(
-            start=start,
-            stop=stop,
-        ) 
+        return NDArray(value=np.arange(*args))
 
     def program_id(self, axis: int):
         if axis == 0:
@@ -284,9 +269,18 @@ def load(self, src: NDArray, *, mask=None, dtype=None, **kwargs):
             value = value.astype(dtype)
         return NDArray(value=value, name=src.name, **kwargs)
 
+    def load_transpose2d(self, src: NDArray, *, mask=None, dtype=None, **kwargs):
+        # THTODO
+        value = src._value
+        return self.load(NDArray(value=value.T, name=src.name), mask=mask, dtype=dtype, **kwargs)
+
     def store(self, dst: NDArray, value: NDArray, *, mask=None, **kwargs):
         if mask is not None:
             value = value[mask]
+        import tp 
+
+        tp.log(f'{dst=}')
+        tp.log(f'{value=}')
         dst._value[:] = value._value[:]
         return dst
 
@@ -300,6 +294,28 @@ def patch():
     nl.arange = lambda *args: nki_builder.arange(*args)
     nl.load = lambda src, **kwargs: nki_builder.load(src, **kwargs)
     nl.store = lambda dst, value, **kwargs: nki_builder.store(dst, value, **kwargs)
+    # see https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/nki/api/nki.language.html
+
+    # TODO: implement
+    # matmul-specific
+    # nl.shared_hbm
+    # nl.psum
+    nl.affine_range
+    nl.par_dim
+    nl.zeros = lambda *args, **kwags: nki_builder.zeros(*args, **kwargs)
+    nl.mgrid
+    #nisa.nc_matmul
+    nl.copy
+
+    # attention-specific
+    nl.load_transpose2d = lambda src, **kwargs: nki_builder.load_transpose2d(src, **kwargs)
+    #nisa.affine_select
+    #nl.tensor_reduce
+    #nisa.activation
+    nl.broadcast_to
+    #nisa.nc_transpose
+
+
 
 
 def unpatch():

From f6b99ae40b06ff9fb2c4eb55f019c0de3fd915ef Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Fri, 26 Sep 2025 04:27:01 -0400
Subject: [PATCH 31/90] remove need for SPMD nki kernel call when tracing

---
 triton_viz/core/trace.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/triton_viz/core/trace.py b/triton_viz/core/trace.py
index e842a134..61745c0b 100644
--- a/triton_viz/core/trace.py
+++ b/triton_viz/core/trace.py
@@ -89,6 +89,9 @@ def __init__(self, kernel, client: str | Client) -> None:
 
     def __getitem__(self, *grid):
         return KernelInterface.__getitem__(self, tuple(*grid))
+    
+    def __call__(self, *args, **kwargs):
+        return self[(1, 1, 1)](*args, **kwargs)
 
 
 #def trace(clients: str | Client | None = None, backend: str = "triton"):

From 4f863fa7a4cc5c10f83c5ab4756438e89c4a4996 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Fri, 26 Sep 2025 04:27:12 -0400
Subject: [PATCH 32/90] remove

---
 examples/nki_softmax.py | 63 -----------------------------------------
 1 file changed, 63 deletions(-)
 delete mode 100644 examples/nki_softmax.py

diff --git a/examples/nki_softmax.py b/examples/nki_softmax.py
deleted file mode 100644
index f84b2711..00000000
--- a/examples/nki_softmax.py
+++ /dev/null
@@ -1,63 +0,0 @@
-from neuronxcc import nki
-import neuronxcc.nki.language as nl
-
-import torch
-import triton_viz
-import torch_xla.core.xla_model as xm
-from triton_viz.clients import Tracer
-from triton_viz.core import config as cfg
-from triton_viz.core.trace import launches
-
-
-@triton_viz.trace(clients=Tracer())
-@nki.jit
-def softmax(in_tensor):
-    # assume B,D; softmax on D for now
-    B, D = in_tensor.shape
-    out_tensor = nl.ndarray((B, D), dtype=in_tensor.dtype, buffer=nl.shared_hbm)
-
-    #assert nl.tile_size.pmax == 128
-    num_tiles = math.ceil(B / 128)
-    for tile_idx in nl.affine_range(num_tiles):
-        i_p = tile_idx * 128 + nl.arange(128)[:, None]
-        i_f = nl.arange(D)[None, :]
-        mask = i_p < B
-        tile = nl.load(in_tensor[i_p, i_f], mask=mask)
-        tile_exp = nl.exp(tile)
-        tile_expsum = nl.sum(tile_exp, -1, keepdims=True, mask=mask)
-        tile_softmax = tile_exp / tile_expsum
-        nl.store(out_tensor[i_p, i_f], value=tile_softmax, mask=mask)
-    return out_tensor
-
-
-if __name__ == "__main__":
-    cfg.reset()
-    device = "cpu"
-    size = 16
-    BLOCK_SIZE = 8
-    torch.manual_seed(0)
-    x = torch.arange(size, dtype=torch.float32, device=device)
-    softmax[(1,1,1)](x)
-
-    # Print records to see what's being captured
-    print(f"Number of launches: {len(launches)}")
-    if launches:
-        launch = launches[-1]
-        print(f"Number of records: {len(launch.records)}")
-        for i, record in enumerate(launch.records):
-            print(f"Record {i}: {type(record).__name__}")
-            if hasattr(record, "ptr"):
-                print(f"  ptr: {record.ptr}")
-            if hasattr(record, "offsets"):
-                print(f"  offsets shape: {record.offsets.shape}")
-            if hasattr(record, "masks"):
-                print(f"  masks shape: {record.masks.shape}")
-
-    # Try to launch visualization
-    try:
-        triton_viz.launch()
-    except Exception as e:
-        print(f"\nError during visualization: {e}")
-        import traceback
-
-        traceback.print_exc()

From 2932bb9138cdc34270bed7a990886f11248162f3 Mon Sep 17 00:00:00 2001
From: gujialiang123 <562960306@qq.com>
Date: Thu, 2 Oct 2025 21:01:32 -0400
Subject: [PATCH 33/90] added calibrate and matmul and show code

---
 examples/matmul_demo.py             |  93 +++++++++++++++++++
 triton_viz/clients/tracer/tracer.py |  55 +++++++++--
 triton_viz/core/data.py             |  17 +++-
 triton_viz/static/gridblock.js      |  64 +++++++++++++
 triton_viz/static/load.js           | 117 ++++++++++++++++++++++-
 triton_viz/static/load_utils.js     |   3 +
 triton_viz/static/matmul.js         |  60 +++++++++++-
 triton_viz/templates/calibrate.html |  96 +++++++++++++++++++
 triton_viz/visualizer/draw.py       |  33 +++++++
 triton_viz/visualizer/interface.py  | 139 ++++++++++++++++++++++++++++
 10 files changed, 657 insertions(+), 20 deletions(-)
 create mode 100644 examples/matmul_demo.py
 create mode 100644 triton_viz/templates/calibrate.html

diff --git a/examples/matmul_demo.py b/examples/matmul_demo.py
new file mode 100644
index 00000000..c6b41db5
--- /dev/null
+++ b/examples/matmul_demo.py
@@ -0,0 +1,93 @@
+import torch
+import triton
+import triton.language as tl
+
+import triton_viz
+from triton_viz.clients import Tracer
+
+
+# Simple matmul kernel producing a C = A @ B (fp32, small sizes for demo)
+@triton_viz.trace(clients=Tracer())
+@triton.jit
+def matmul_kernel(
+    a_ptr,
+    b_ptr,
+    c_ptr,
+    M: tl.constexpr,
+    N: tl.constexpr,
+    K: tl.constexpr,
+    stride_am,
+    stride_ak,
+    stride_bk,
+    stride_bn,
+    stride_cm,
+    stride_cn,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+    BLOCK_K: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+    offs_k = tl.arange(0, BLOCK_K)
+
+    a_ptrs = a_ptr + offs_m[:, None] * stride_am + offs_k[None, :] * stride_ak
+    b_ptrs = b_ptr + offs_k[:, None] * stride_bk + offs_n[None, :] * stride_bn
+
+    acc = tl.zeros((BLOCK_M, BLOCK_N), dtype=tl.float32)
+
+    for k in range(0, K, BLOCK_K):
+        a = tl.load(
+            a_ptrs, mask=(offs_m[:, None] < M) & (k + offs_k[None, :] < K), other=0.0
+        )
+        b = tl.load(
+            b_ptrs, mask=(k + offs_k[:, None] < K) & (offs_n[None, :] < N), other=0.0
+        )
+        acc += tl.dot(a, b)
+        a_ptrs += BLOCK_K * stride_ak
+        b_ptrs += BLOCK_K * stride_bk
+
+    c_ptrs = c_ptr + offs_m[:, None] * stride_cm + offs_n[None, :] * stride_cn
+    tl.store(c_ptrs, acc, mask=(offs_m[:, None] < M) & (offs_n[None, :] < N))
+
+
+def run_demo():
+    torch.manual_seed(0)
+    M, N, K = 64, 64, 64
+    a = torch.randn((M, K), dtype=torch.float32)
+    b = torch.randn((K, N), dtype=torch.float32)
+    c = torch.empty((M, N), dtype=torch.float32)
+
+    BLOCK_M, BLOCK_N, BLOCK_K = 32, 32, 32
+    grid = (triton.cdiv(M, BLOCK_M), triton.cdiv(N, BLOCK_N))
+
+    matmul_kernel[grid](
+        a,
+        b,
+        c,
+        M,
+        N,
+        K,
+        a.stride(0),
+        a.stride(1),
+        b.stride(0),
+        b.stride(1),
+        c.stride(0),
+        c.stride(1),
+        BLOCK_M=BLOCK_M,
+        BLOCK_N=BLOCK_N,
+        BLOCK_K=BLOCK_K,
+    )
+
+    # Verify correctness
+    ref = a @ b
+    assert torch.allclose(c, ref, atol=1e-3), "matmul result mismatch"
+
+    # Launch viz UI in blocking (share=True) mode
+    triton_viz.launch(share=True)
+
+
+if __name__ == "__main__":
+    run_demo()
diff --git a/triton_viz/clients/tracer/tracer.py b/triton_viz/clients/tracer/tracer.py
index 4b924059..323cd96d 100644
--- a/triton_viz/clients/tracer/tracer.py
+++ b/triton_viz/clients/tracer/tracer.py
@@ -3,6 +3,7 @@
 from ...core.data import Op, Load, Store, ReduceSum, Dot, Grid, RawLoad, RawStore
 from typing import Callable, Optional, Union
 import numpy as np
+import traceback
 
 
 def _convert_grid_idx(grid_idx) -> Optional[tuple[int, int, int]]:
@@ -64,6 +65,35 @@ def grid_callback(self, grid: tuple[int, ...]):
         self.tensors = sorted(self.tensors, key=lambda x: x.data_ptr())
 
     def register_op_callback(self, op_type: type[Op]) -> OpCallbacks:
+        def _extract_user_frames() -> list[traceback.FrameSummary]:
+            stack: list[traceback.FrameSummary] = list(traceback.extract_stack())
+            # drop current frames (this function and callers)
+            stack = stack[:-2]
+            cleaned: list[traceback.FrameSummary] = []
+            for f in stack:
+                fn = f.filename.replace("\\", "/")
+                if any(
+                    s in fn
+                    for s in [
+                        "triton_viz/core/",
+                        "triton_viz/clients/",
+                        "triton/runtime/",
+                        "triton/language/",
+                        "site-packages/triton/",
+                        "runpy.py",
+                        "IPython",
+                    ]
+                ):
+                    continue
+                cleaned.append(f)
+            if cleaned:
+                return cleaned
+            # fallback to last non "<...>" frame
+            for f in reversed(stack):
+                if not f.filename.startswith("<"):
+                    return [f]
+            return stack[-1:]
+
         def pre_load_callback(
             ptr, mask, other, cache_modifier, eviction_policy, is_volatile
         ):
@@ -71,18 +101,18 @@ def pre_load_callback(
                 return
             first_ptr = np.reshape(ptr.data, (-1))[0]
             tensor = self._get_tensor(first_ptr)
-            self.records.append(
-                Load(tensor.data_ptr(), ptr.data - tensor.data_ptr(), mask.data)
-            )
+            rec = Load(tensor.data_ptr(), ptr.data - tensor.data_ptr(), mask.data)
+            rec.call_path = _extract_user_frames()
+            self.records.append(rec)
 
         def pre_store_callback(ptr, value, mask, cache_modifier, eviction_policy):
             if not self.sample:
                 return
             first_ptr = np.reshape(ptr.data, (-1))[0]
             tensor = self._get_tensor(first_ptr)
-            self.records.append(
-                Store(tensor.data_ptr(), ptr.data - tensor.data_ptr(), mask.data)
-            )
+            rec = Store(tensor.data_ptr(), ptr.data - tensor.data_ptr(), mask.data)
+            rec.call_path = _extract_user_frames()
+            self.records.append(rec)
 
         # Raw (unmasked) ops: synthesize a full True mask based on ptr shape
         def pre_raw_load_callback(ptr):
@@ -92,7 +122,9 @@ def pre_raw_load_callback(ptr):
             tensor = self._get_tensor(first_ptr)
             offsets = ptr.data - tensor.data_ptr()
             true_mask = np.ones_like(offsets, dtype=bool)
-            self.records.append(Load(tensor.data_ptr(), offsets, true_mask))
+            rec = Load(tensor.data_ptr(), offsets, true_mask)
+            rec.call_path = _extract_user_frames()
+            self.records.append(rec)
 
         def pre_raw_store_callback(ptr, value):
             if not self.sample:
@@ -101,7 +133,9 @@ def pre_raw_store_callback(ptr, value):
             tensor = self._get_tensor(first_ptr)
             offsets = ptr.data - tensor.data_ptr()
             true_mask = np.ones_like(offsets, dtype=bool)
-            self.records.append(Store(tensor.data_ptr(), offsets, true_mask))
+            rec = Store(tensor.data_ptr(), offsets, true_mask)
+            rec.call_path = _extract_user_frames()
+            self.records.append(rec)
 
         def post_reduce_sum_callback(ret, input, axis=None, keep_dims=False):
             if not self.sample:
@@ -116,7 +150,10 @@ def post_dot_callback(ret, input, other, *args):
             input_shape = input.data.shape
             other_shape = other.data.shape
             ret_shape = ret.data.shape
-            self.records.append(Dot(input_shape, other_shape, ret_shape))
+            # Pass input/other raw arrays so draw.py can render MatMul
+            rec = Dot(input_shape, other_shape, ret_shape, input.data, other.data)
+            rec.call_path = _extract_user_frames()
+            self.records.append(rec)
 
         if op_type is Load:
             return OpCallbacks(before_callback=pre_load_callback)
diff --git a/triton_viz/core/data.py b/triton_viz/core/data.py
index 41182048..7f259e58 100644
--- a/triton_viz/core/data.py
+++ b/triton_viz/core/data.py
@@ -16,17 +16,26 @@
 @dataclass
 class Op:
     name: ClassVar[str] = "op"
-    call_path: list[traceback.StackSummary] = field(init=False, default_factory=list)
+    call_path: list[traceback.FrameSummary] = field(init=False, default_factory=list)
 
     def __post_init__(self):
-        self.call_path = traceback.extract_stack()[:-2]
+        full_stack = traceback.extract_stack()[:-2]
+        # keep original for fallback
+        self.call_path = full_stack
         clean_call_path = []
-        for frame in self.call_path:
+        for frame in full_stack:
             if not any(
                 triton_frame in frame.filename for triton_frame in TRITON_FRAMES
             ):
                 clean_call_path.append(frame)
-        self.call_path = clean_call_path
+        # if filtering removed all frames, fallback to last meaningful frame(s)
+        if clean_call_path:
+            self.call_path = clean_call_path
+        else:
+            for frame in reversed(full_stack):
+                if not str(frame.filename).startswith("<"):
+                    self.call_path = [frame]
+                    break
 
 
 @dataclass
diff --git a/triton_viz/static/gridblock.js b/triton_viz/static/gridblock.js
index bbdda5bb..7ee9762f 100644
--- a/triton_viz/static/gridblock.js
+++ b/triton_viz/static/gridblock.js
@@ -90,6 +90,70 @@ export class GridBlock {
         if (this.blockData.length > 0) {
             this.displayOpVisualization(this.blockData[0]);
         }
+
+        // Add a global "Show Code" toggle once (applies to any op type)
+        const codeBtnId = 'global-code-toggle-btn';
+        if (!document.getElementById(codeBtnId)) {
+            const btn = document.createElement('button');
+            btn.id = codeBtnId;
+            btn.textContent = 'Show Code: OFF';
+            Object.assign(btn.style, {
+                position: 'fixed',
+                right: '10px',
+                top: '10px',
+                zIndex: '2001'
+            });
+            document.body.appendChild(btn);
+
+            let panel = null;
+            const destroyPanel = () => { if (panel && panel.remove) panel.remove(); panel = null; };
+            const createPanel = async (uuid, frameIdx = 0, context = 8) => {
+                destroyPanel();
+                const wrapper = document.createElement('div');
+                Object.assign(wrapper.style, {
+                    position: 'fixed', right: '10px', top: '50px', width: '520px', maxHeight: '60vh', overflow: 'auto',
+                    padding: '8px 10px', background: 'rgba(0,0,0,0.65)', color: '#fff', font: '12px Menlo, Consolas, monospace',
+                    borderRadius: '6px', zIndex: '2000'
+                });
+                const header = document.createElement('div');
+                header.textContent = 'Operation Code & Context';
+                header.style.marginBottom = '6px';
+                header.style.opacity = '0.9';
+                wrapper.appendChild(header);
+                try {
+                    const res = await fetch('/api/op_code', { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ uuid, frame_idx: frameIdx, context }) });
+                    const data = await res.json();
+                    const meta = document.createElement('div');
+                    meta.style.marginBottom = '4px';
+                    meta.textContent = `${data.filename || ''}:${data.lineno || ''}`;
+                    wrapper.appendChild(meta);
+                    const pre = document.createElement('pre');
+                    pre.style.margin = '0';
+                    pre.style.whiteSpace = 'pre';
+                    const lines = (data.lines || []).map(l => {
+                        const mark = (data.highlight === l.no) ? '▶ ' : '  ';
+                        return `${mark}${String(l.no).padStart(6,' ')} | ${l.text||''}`;
+                    }).join('\n');
+                    pre.textContent = lines || '(no code available)';
+                    wrapper.appendChild(pre);
+                } catch (e) {
+                    const err = document.createElement('div');
+                    err.textContent = 'Failed to load code context.';
+                    wrapper.appendChild(err);
+                }
+                document.body.appendChild(wrapper);
+                panel = wrapper;
+            };
+
+            btn.addEventListener('click', async () => {
+                const turnOn = btn.textContent.endsWith('OFF');
+                btn.textContent = `Show Code: ${turnOn ? 'ON' : 'OFF'}`;
+                if (!this.blockData || this.blockData.length === 0) { if (!turnOn) destroyPanel(); return; }
+                // Prefer latest clicked/active uuid; fallback to first in current grid block
+                const activeUuid = window.current_op_uuid || (this.blockData[0] && this.blockData[0].uuid);
+                if (turnOn && activeUuid) await createPanel(activeUuid, 0, 8); else destroyPanel();
+            });
+        }
     }
 
     createVisualizationContainer() {
diff --git a/triton_viz/static/load.js b/triton_viz/static/load.js
index 73a86e8f..25274c3f 100644
--- a/triton_viz/static/load.js
+++ b/triton_viz/static/load.js
@@ -25,6 +25,9 @@ export function createLoadVisualization(containerElement, op) {
         .then(data => console.log('Set current op:', data))
         .catch((error) => console.error('Error:', error));
 
+        // expose current op uuid globally for generic code panel in gridblock
+        try { window.current_op_uuid = op.uuid; } catch(e){}
+
         let currentStep = 0;
         let frame = 0;
         let isPaused = false;
@@ -63,6 +66,28 @@ export function createLoadVisualization(containerElement, op) {
         const dragToggle = document.createElement('button');
         dragToggle.textContent = 'Drag Cubes: OFF';
         controlBar.appendChild(dragToggle);
+        const codeToggle = document.createElement('button');
+        codeToggle.textContent = 'Show Code: OFF';
+        controlBar.appendChild(codeToggle);
+        // Mouse pick calibration (dx/dy in pixels)
+        const calibWrap = document.createElement('div');
+        calibWrap.style.display = 'flex';
+        calibWrap.style.alignItems = 'center';
+        calibWrap.style.gap = '4px';
+        const calibLabel = document.createElement('span');
+        calibLabel.textContent = 'Calib:';
+        calibLabel.style.opacity = '0.8';
+        const dxMinus = document.createElement('button'); dxMinus.textContent = '−X';
+        const dxPlus  = document.createElement('button'); dxPlus.textContent  = '+X';
+        const dyMinus = document.createElement('button'); dyMinus.textContent = '−Y';
+        const dyPlus  = document.createElement('button'); dyPlus.textContent  = '+Y';
+        const dxdyInfo = document.createElement('span'); dxdyInfo.style.minWidth = '70px'; dxdyInfo.style.textAlign = 'center';
+        const dxdyReset = document.createElement('button'); dxdyReset.textContent = 'Reset';
+        calibWrap.appendChild(calibLabel);
+        calibWrap.appendChild(dxMinus); calibWrap.appendChild(dxPlus);
+        calibWrap.appendChild(dyMinus); calibWrap.appendChild(dyPlus);
+        calibWrap.appendChild(dxdyInfo); calibWrap.appendChild(dxdyReset);
+        controlBar.appendChild(calibWrap);
         containerElement.appendChild(controlBar);
 
         // expose for debugging
@@ -79,6 +104,7 @@ export function createLoadVisualization(containerElement, op) {
         let legendEl = null;
         let scheme = 'mono';
         let monoBaseHex = '#3b82f6';
+        let codePanel = null;
 
         const COLOR_GLOBAL = new THREE.Color(0.2, 0.2, 0.2);    // Dark Gray
         const COLOR_SLICE = new THREE.Color(0.0, 0.7, 1.0);     // Cyan (starting color for global slice)
@@ -116,6 +142,11 @@ export function createLoadVisualization(containerElement, op) {
 
         const raycaster = new THREE.Raycaster();
         const mouse = new THREE.Vector2();
+        // persistent calibration offsets
+        let mouseDx = Number(localStorage.getItem('viz_mouse_dx') || 0);
+        let mouseDy = Number(localStorage.getItem('viz_mouse_dy') || 0);
+        function updateDxDyLabel(){ dxdyInfo.textContent = `dx=${mouseDx}, dy=${mouseDy}`; }
+        updateDxDyLabel();
         // Drag state
         let dragModeOn = false;
         let isDragging = false;
@@ -135,8 +166,13 @@ export function createLoadVisualization(containerElement, op) {
         animate();
 
         function _updateMouseNDC(event) {
-            mouse.x = (event.clientX / containerElement.clientWidth) * 2 - 1;
-            mouse.y = -(event.clientY / containerElement.clientHeight) * 2 + 1;
+            const rect = renderer.domElement.getBoundingClientRect();
+            const dpr = (window.devicePixelRatio || 1);
+            const px = (event.clientX - rect.left + mouseDx) * dpr;
+            const py = (event.clientY - rect.top  + mouseDy) * dpr;
+            const w = rect.width * dpr, h = rect.height * dpr;
+            mouse.x = (px / w) * 2 - 1;
+            mouse.y = -(py / h) * 2 + 1;
         }
 
         function _raycastAll() {
@@ -157,8 +193,7 @@ export function createLoadVisualization(containerElement, op) {
         }
 
         async function onMouseMove(event) {
-            mouse.x = (event.clientX / containerElement.clientWidth) * 2 - 1;
-            mouse.y = -(event.clientY / containerElement.clientHeight) * 2 + 1;
+            _updateMouseNDC(event);
 
             raycaster.setFromCamera(mouse, camera);
 
@@ -288,6 +323,63 @@ export function createLoadVisualization(containerElement, op) {
             legendEl = wrapper;
         }
 
+        function destroyCodePanel() {
+            if (codePanel && codePanel.remove) codePanel.remove();
+            codePanel = null;
+        }
+
+        async function createCodePanel(frameIdx = 0, context = 8) {
+            destroyCodePanel();
+            const wrapper = document.createElement('div');
+            wrapper.style.position = 'fixed';
+            wrapper.style.right = '10px';
+            wrapper.style.top = '50px';
+            wrapper.style.width = '520px';
+            wrapper.style.maxHeight = '60vh';
+            wrapper.style.overflow = 'auto';
+            wrapper.style.padding = '8px 10px';
+            wrapper.style.background = 'rgba(0,0,0,0.65)';
+            wrapper.style.color = '#fff';
+            wrapper.style.font = '12px Menlo, Consolas, monospace';
+            wrapper.style.borderRadius = '6px';
+            wrapper.style.zIndex = '2000';
+
+            const header = document.createElement('div');
+            header.textContent = 'Operation Code & Context';
+            header.style.marginBottom = '6px';
+            header.style.opacity = '0.9';
+            wrapper.appendChild(header);
+
+            try {
+                const res = await fetch('/api/op_code', {
+                    method: 'POST',
+                    headers: { 'Content-Type': 'application/json' },
+                    body: JSON.stringify({ uuid: op.uuid, frame_idx: frameIdx, context })
+                });
+                const data = await res.json();
+                const meta = document.createElement('div');
+                meta.style.marginBottom = '4px';
+                meta.textContent = `${data.filename || ''}:${data.lineno || ''}`;
+                wrapper.appendChild(meta);
+                const pre = document.createElement('pre');
+                pre.style.margin = '0';
+                pre.style.whiteSpace = 'pre';
+                const lines = (data.lines || []).map(l => {
+                    const mark = (data.highlight === l.no) ? '▶ ' : '  ';
+                    return `${mark}${String(l.no).padStart(6,' ')} | ${l.text||''}`;
+                }).join('\n');
+                pre.textContent = lines || '(no code available)';
+                wrapper.appendChild(pre);
+            } catch (e) {
+                const err = document.createElement('div');
+                err.textContent = 'Failed to load code context.';
+                wrapper.appendChild(err);
+            }
+
+            containerElement.appendChild(wrapper);
+            codePanel = wrapper;
+        }
+
         function applyColorMapIfNeeded() {
             if (!colorizeOn || !tensorCache) return;
             const { min, max, dims, values } = tensorCache;
@@ -480,6 +572,23 @@ export function createLoadVisualization(containerElement, op) {
             orbitControls.enabled = !dragModeOn;
         });
 
+        codeToggle.addEventListener('click', async () => {
+            const on = codeToggle.textContent.endsWith('OFF');
+            codeToggle.textContent = `Show Code: ${on ? 'ON' : 'OFF'}`;
+            if (on) {
+                await createCodePanel(0, 8);
+            } else {
+                destroyCodePanel();
+            }
+        });
+
+        // calibration handlers
+        dxMinus.addEventListener('click', ()=>{ mouseDx -= 1; localStorage.setItem('viz_mouse_dx', String(mouseDx)); updateDxDyLabel(); });
+        dxPlus.addEventListener('click',  ()=>{ mouseDx += 1; localStorage.setItem('viz_mouse_dx', String(mouseDx)); updateDxDyLabel(); });
+        dyMinus.addEventListener('click', ()=>{ mouseDy -= 1; localStorage.setItem('viz_mouse_dy', String(mouseDy)); updateDxDyLabel(); });
+        dyPlus.addEventListener('click',  ()=>{ mouseDy += 1; localStorage.setItem('viz_mouse_dy', String(mouseDy)); updateDxDyLabel(); });
+        dxdyReset.addEventListener('click', ()=>{ mouseDx = 0; mouseDy = 0; localStorage.setItem('viz_mouse_dx','0'); localStorage.setItem('viz_mouse_dy','0'); updateDxDyLabel(); });
+
         function updateSideMenu(tensorName, x, y, z, value) {
             if (!tensorName) {
                 sideMenu.innerHTML = '';
diff --git a/triton_viz/static/load_utils.js b/triton_viz/static/load_utils.js
index 52ea6ff9..5a38bc70 100644
--- a/triton_viz/static/load_utils.js
+++ b/triton_viz/static/load_utils.js
@@ -12,6 +12,9 @@ export function setupScene(container, backgroundColor = 0x000000) {
     scene.background = new THREE.Color(backgroundColor);
     const camera = new THREE.PerspectiveCamera(45, container.clientWidth / container.clientHeight, 0.1, 1000);
     const renderer = new THREE.WebGLRenderer({ antialias: true });
+    // Honour device pixel ratio to align raycaster with drawn pixels
+    const dpr = (window.devicePixelRatio || 1);
+    renderer.setPixelRatio(dpr);
     renderer.setSize(container.clientWidth, container.clientHeight);
     container.appendChild(renderer.domElement);
 
diff --git a/triton_viz/static/matmul.js b/triton_viz/static/matmul.js
index c32e9285..3ce14757 100644
--- a/triton_viz/static/matmul.js
+++ b/triton_viz/static/matmul.js
@@ -1,4 +1,5 @@
 import * as THREE from 'https://esm.sh/three@0.155.0/build/three.module.js';
+import { OrbitControls } from 'https://esm.sh/three@0.155.0/examples/jsm/controls/OrbitControls.js';
 
 export function createMatMulVisualization(containerElement, op) {
     const { input_shape, other_shape, output_shape } = op;
@@ -80,11 +81,23 @@ export function createMatMulVisualization(containerElement, op) {
 
     const raycaster = new THREE.Raycaster();
     const mouse = new THREE.Vector2();
-
+    // persistent calibration offsets shared with Load view for consistency
+    let mouseDx = Number(localStorage.getItem('viz_mouse_dx') || 0);
+    let mouseDy = Number(localStorage.getItem('viz_mouse_dy') || 0);
+
+
+    function _updateMouseNDC(event) {
+        const rect = renderer.domElement.getBoundingClientRect();
+        const dpr = (window.devicePixelRatio || 1);
+        const px = (event.clientX - rect.left + mouseDx) * dpr;
+        const py = (event.clientY - rect.top  + mouseDy) * dpr;
+        const w = rect.width * dpr, h = rect.height * dpr;
+        mouse.x = (px / w) * 2 - 1;
+        mouse.y = -(py / h) * 2 + 1;
+    }
 
     async function onMouseMove(event) {
-        mouse.x = (event.clientX / containerElement.clientWidth) * 2 - 1;
-        mouse.y = -(event.clientY / containerElement.clientHeight) * 2 + 1;
+        _updateMouseNDC(event);
 
         raycaster.setFromCamera(mouse, camera);
 
@@ -150,6 +163,8 @@ export function createMatMulVisualization(containerElement, op) {
     const camera = new THREE.PerspectiveCamera(45, containerElement.clientWidth / containerElement.clientHeight, 0.1, 1000);
 
     const renderer = new THREE.WebGLRenderer({ antialias: true });
+    const dpr = (window.devicePixelRatio || 1);
+    renderer.setPixelRatio(dpr);
     renderer.setSize(containerElement.clientWidth, containerElement.clientHeight);
     containerElement.appendChild(renderer.domElement);
 
@@ -221,6 +236,11 @@ export function createMatMulVisualization(containerElement, op) {
 
     camera.position.set(center.x, center.y, center.z + cameraZ);
     camera.lookAt(center);
+    const controls = new OrbitControls(camera, renderer.domElement);
+    controls.enableDamping = true;
+    controls.dampingFactor = 0.05;
+    controls.target.copy(center);
+    controls.update();
 
     let isPaused = false;
 
@@ -241,10 +261,12 @@ export function createMatMulVisualization(containerElement, op) {
     function resetColors() {
         matrixA.children.forEach(cube => cube.material.color.copy(COLOR_A));
         matrixB.children.forEach(cube => cube.material.color.copy(COLOR_B));
+        matrixC.children.forEach(cube => cube.material.color.copy(COLOR_C));
     }
 
     function animate() {
         requestAnimationFrame(animate);
+        controls.update();
 
         if (!isPaused && frame < totalFrames) {
             resetColors();
@@ -296,6 +318,37 @@ export function createMatMulVisualization(containerElement, op) {
     controlPanel.appendChild(playPauseButton);
     controlPanel.appendChild(resetButton);
 
+    // Color by Value (for C matrix) toggle + legend (mono blue)
+    const colorToggle = document.createElement('button');
+    colorToggle.textContent = 'Color by Value: OFF';
+    controlPanel.appendChild(colorToggle);
+    let colorOn = false;
+    let legendEl = null;
+    function destroyLegend(){ if(legendEl && legendEl.remove) legendEl.remove(); legendEl=null; }
+    function createLegend(min,max){
+        destroyLegend();
+        const w = document.createElement('div');
+        Object.assign(w.style,{position:'absolute', left:'10px', bottom:'60px', background:'rgba(0,0,0,0.6)', color:'#fff', padding:'6px 8px', borderRadius:'6px'});
+        const c = document.createElement('canvas'); c.width=220; c.height=10; const ctx=c.getContext('2d');
+        for(let x=0;x<c.width;x++){ const t=x/(c.width-1); const r=t,g=0.2,b=1-t; ctx.fillStyle=`rgb(${Math.round(r*255)},${Math.round(g*255)},${Math.round(b*255)})`; ctx.fillRect(x,0,1,c.height);}
+        const lab=document.createElement('div'); lab.style.display='flex'; lab.style.justifyContent='space-between'; lab.style.marginTop='2px'; lab.innerHTML=`<span>${min.toFixed(3)}</span><span>${max.toFixed(3)}</span>`;
+        const ttl=document.createElement('div'); ttl.textContent='Value (C)'; ttl.style.marginBottom='4px'; ttl.style.opacity='0.9';
+        w.appendChild(ttl); w.appendChild(c); w.appendChild(lab); containerElement.appendChild(w); legendEl=w;
+    }
+    async function fetchCValues(){
+        try{ const res=await fetch('/api/getMatmulC',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({uuid: op.uuid})});
+            return await res.json(); }catch(e){ return {error:String(e)} }
+    }
+    function applyColorCWithData(data){ if(!colorOn||!data||data.error) return; try{
+            const M=(data.shape||[])[0]||0, N=(data.shape||[])[1]||0; const vals2d=data.values||[]; const vals=[]; const cells=[];
+            matrixC.children.forEach((cube,i)=>{ const row=Math.floor(i/N), col=i%N; cells.push({cube,row,col}); vals.push((row<M&&col<N)? vals2d[row][col]:0.0); });
+            const mn=Math.min(...vals), mx=Math.max(...vals);
+            cells.forEach((e,idx)=>{ const v=vals[idx]; const u=(mx===mn)?0.5:(v-mn)/(mx-mn); const r=u,g=0.2,b=1-u; e.cube.material.color.setRGB(r,g,b); });
+            createLegend(mn,mx);
+        }catch(e){}
+    }
+    colorToggle.addEventListener('click', async ()=>{ colorOn=!colorOn; colorToggle.textContent=`Color by Value: ${colorOn?'ON':'OFF'}`; if(!colorOn){ destroyLegend(); resetColors(); } else { const data=await fetchCValues(); applyColorCWithData(data); } });
+
     containerElement.appendChild(controlPanel);
 
     const PAN_SPEED = 0.1;
@@ -343,6 +396,7 @@ export function createMatMulVisualization(containerElement, op) {
     window.addEventListener('resize', onResize);
     window.addEventListener('keydown', onKeyDown);
     containerElement.addEventListener('mousemove', onMouseMove);
+    try { window.current_op_uuid = op.uuid; } catch (e) {}
 
     // Mouse wheel zoom for matmul view
     const WHEEL_ZOOM_SPEED = 0.5;
diff --git a/triton_viz/templates/calibrate.html b/triton_viz/templates/calibrate.html
new file mode 100644
index 00000000..3e3d20da
--- /dev/null
+++ b/triton_viz/templates/calibrate.html
@@ -0,0 +1,96 @@
+<!DOCTYPE html>
+<html lang="en">
+<head>
+    <meta charset="UTF-8" />
+    <meta name="viewport" content="width=device-width, initial-scale=1.0" />
+    <title>Mouse Calibration</title>
+    <style>
+        body { margin:0; background:#111; color:#ddd; font-family: Arial, sans-serif; }
+        #panel { position: fixed; top: 10px; left: 10px; background: rgba(0,0,0,0.6); padding: 8px 10px; border-radius: 6px; }
+        button { margin-right: 6px; }
+        #info { margin-top: 6px; opacity: 0.9; }
+        canvas { display:block; margin: 0 auto; outline: 1px solid #333; }
+    </style>
+</head>
+<body>
+    <div id="panel">
+        <div>点击每个目标中心点进行采样（至少 4 次）。</div>
+        <div id="info">样本: 0 | dx=0, dy=0</div>
+        <button id="reset">重置</button>
+        <button id="apply">应用校准</button>
+    </div>
+    <canvas id="cv" width="800" height="500"></canvas>
+    <script type="module">
+    const cv = document.getElementById('cv');
+    const ctx = cv.getContext('2d');
+    const info = document.getElementById('info');
+    const resetBtn = document.getElementById('reset');
+    const applyBtn = document.getElementById('apply');
+
+    const targets = [
+        {x:150,y:120},{x:400,y:120},{x:650,y:120},
+        {x:150,y:360},{x:400,y:360},{x:650,y:360}
+    ];
+    let samples = [];
+
+    function draw(){
+        ctx.clearRect(0,0,cv.width,cv.height);
+        // grid
+        ctx.strokeStyle = '#222';
+        for(let x=0;x<cv.width;x+=20){ ctx.beginPath(); ctx.moveTo(x,0); ctx.lineTo(x,cv.height); ctx.stroke(); }
+        for(let y=0;y<cv.height;y+=20){ ctx.beginPath(); ctx.moveTo(0,y); ctx.lineTo(cv.width,y); ctx.stroke(); }
+        // targets
+        targets.forEach(t=>{
+            ctx.fillStyle = '#0ff';
+            ctx.fillRect(t.x-6,t.y-6,12,12);
+            ctx.strokeStyle = '#08f';
+            ctx.strokeRect(t.x-10,t.y-10,20,20);
+        });
+        // clicks
+        samples.forEach(s=>{
+            ctx.fillStyle = '#ff0';
+            ctx.beginPath(); ctx.arc(s.px,s.py,4,0,Math.PI*2); ctx.fill();
+        });
+    }
+
+    function computeDxDy(){
+        if(samples.length === 0) return {dx:0,dy:0};
+        let sumDx=0,sumDy=0,n=0;
+        for(const s of samples){
+            sumDx += (s.tx - s.px);
+            sumDy += (s.ty - s.py);
+            n++;
+        }
+        return {dx: Math.round(sumDx/n), dy: Math.round(sumDy/n)};
+    }
+
+    cv.addEventListener('click', (e)=>{
+        const rect = cv.getBoundingClientRect();
+        const px = e.clientX - rect.left;
+        const py = e.clientY - rect.top;
+        // pick nearest target
+        let best=null, bestd=1e9;
+        for(const t of targets){
+            const d=(t.x-px)**2+(t.y-py)**2; if(d<bestd){best={tx:t.x,ty:t.y}; bestd=d;}
+        }
+        samples.push({px,py, ...best});
+        const {dx,dy} = computeDxDy();
+        info.textContent = `样本: ${samples.length} | dx=${dx}, dy=${dy}`;
+        draw();
+    });
+
+    resetBtn.addEventListener('click', ()=>{
+        samples = []; info.textContent = '样本: 0 | dx=0, dy=0'; draw();
+    });
+
+    applyBtn.addEventListener('click', ()=>{
+        const {dx,dy} = computeDxDy();
+        localStorage.setItem('viz_mouse_dx', String(dx));
+        localStorage.setItem('viz_mouse_dy', String(dy));
+        alert(`已应用校准: dx=${dx}, dy=${dy}`);
+    });
+
+    draw();
+    </script>
+    </body>
+    </html>
diff --git a/triton_viz/visualizer/draw.py b/triton_viz/visualizer/draw.py
index 23feb087..22dd6a1b 100644
--- a/triton_viz/visualizer/draw.py
+++ b/triton_viz/visualizer/draw.py
@@ -169,6 +169,8 @@ def prepare_visualization_data(program_records, tensor_table):
                     "other_shape": record.other_shape,
                     "output_shape": record.output_shape,
                     "uuid": record_uuid,
+                    # provide C values for color-by-value
+                    "c_shape": record.output_shape,
                 }
             )
 
@@ -176,6 +178,16 @@ def prepare_visualization_data(program_records, tensor_table):
                 "input_data": torch.tensor(record.input_data),
                 "other_data": torch.tensor(record.other_data),
                 "intermediate_results": record.intermediate_results,
+                "tracebacks": [
+                    {
+                        "filename": f.filename,
+                        "lineno": f.lineno,
+                        "line": f.line,
+                        "name": f.name,
+                    }
+                    for f in getattr(record, "call_path", [])
+                ],
+                # prepare C values after kernel (if available from intermediate or recompute best-effort)
             }
 
         elif isinstance(record, Load):
@@ -197,6 +209,15 @@ def prepare_visualization_data(program_records, tensor_table):
             raw_tensor_data[record_uuid] = {
                 "global_tensor": global_tensor.data.cpu(),  # Ensure it's on CPU
                 "dims": len(global_tensor.data.cpu().shape),
+                "tracebacks": [
+                    {
+                        "filename": f.filename,
+                        "lineno": f.lineno,
+                        "line": f.line,
+                        "name": f.name,
+                    }
+                    for f in getattr(record, "call_path", [])
+                ],
             }
             print(record.masks.shape)
 
@@ -216,6 +237,18 @@ def prepare_visualization_data(program_records, tensor_table):
                 }
             )
 
+            raw_tensor_data[record_uuid] = {
+                "tracebacks": [
+                    {
+                        "filename": f.filename,
+                        "lineno": f.lineno,
+                        "line": f.line,
+                        "name": f.name,
+                    }
+                    for f in getattr(record, "call_path", [])
+                ],
+            }
+
     return visualization_data, raw_tensor_data, ""
 
 
diff --git a/triton_viz/visualizer/interface.py b/triton_viz/visualizer/interface.py
index 9fa30b5d..95b4f0de 100644
--- a/triton_viz/visualizer/interface.py
+++ b/triton_viz/visualizer/interface.py
@@ -113,6 +113,35 @@ def update_global_data():
     global_data["analysis"] = df_dict
 
 
+def _safe_read_file_segment(filename: str, lineno: int, context: int = 8):
+    try:
+        # only allow files under current working dir for safety
+        cwd = os.path.realpath(os.getcwd())
+        path = os.path.realpath(filename)
+        if not path.startswith(cwd):
+            return None
+        start = max(1, lineno - context)
+        end = lineno + context
+        lines = []
+        with open(path, "r", encoding="utf-8", errors="ignore") as f:
+            for i, line in enumerate(f, start=1):
+                if i < start:
+                    continue
+                if i > end:
+                    break
+                lines.append({"no": i, "text": line.rstrip("\n")})
+        return {
+            "filename": path,
+            "lineno": lineno,
+            "start": start,
+            "end": end,
+            "highlight": lineno,
+            "lines": lines,
+        }
+    except Exception:
+        return None
+
+
 @app.route("/")
 def index():
     update_global_data()
@@ -125,6 +154,12 @@ def debug_page():
     return render_template("debug.html")
 
 
+@app.route("/calibrate")
+def calibrate_page():
+    # A minimal page to calibrate mouse picking dx/dy
+    return render_template("calibrate.html")
+
+
 @app.route("/api/data")
 def get_data():
     global global_data
@@ -148,6 +183,110 @@ def set_current_op():
     )
 
 
+@app.route("/api/op_code", methods=["POST"])
+def get_op_code():
+    global raw_tensor_data, current_fullscreen_op
+    data = request.json or {}
+    # ensure data prepared
+    if raw_tensor_data is None:
+        update_global_data()
+    uuid = data.get("uuid") or current_fullscreen_op
+    frame_idx = int(data.get("frame_idx", 0))
+    context = int(data.get("context", 8))
+    if not uuid or raw_tensor_data is None or uuid not in raw_tensor_data:
+        return jsonify({"error": "Operation not found"}), 404
+    tb_list = raw_tensor_data[uuid].get("tracebacks") or []
+    if not tb_list:
+        return jsonify({"error": "Traceback not available"}), 200
+    # Heuristic: pick the BEST user frame under CWD (closest to op)
+    cwd = os.path.realpath(os.getcwd())
+
+    def _score(tb: dict) -> int:
+        fn = tb.get("filename") or ""
+        line = tb.get("line") or ""
+        name = tb.get("name") or ""
+        p = os.path.realpath(fn)
+        score = 0
+        if p.startswith(cwd):
+            score += 5
+        if any(
+            s in p
+            for s in ["site-packages", "triton_viz/", "triton/", "runpy.py", "IPython"]
+        ):
+            score -= 10
+        if name.endswith("_kernel") or "kernel" in name:
+            score += 3
+        if "tl." in line:
+            score += 2
+        if "examples" in p:
+            score += 1
+        return score
+
+    # Prefer frames from tail (closest to current) and highest score
+    best = None
+    best_score = -(10**9)
+    for tb in reversed(tb_list):
+        sc = _score(tb)
+        if sc > best_score:
+            best, best_score = tb, sc
+    chosen = best if best is not None else None
+    if chosen is None:
+        # fallback: use requested frame or last non-<string> frame
+        frame_idx = max(0, min(frame_idx, len(tb_list) - 1))
+        chosen = tb_list[frame_idx]
+        for tb in reversed(tb_list):
+            fn = tb.get("filename") or ""
+            if not fn.startswith("<"):
+                chosen = tb
+                break
+    tb = chosen
+    filename = tb.get("filename")
+    lineno = int(tb.get("lineno", 0))
+    line_of_code = tb.get("line")
+    seg = _safe_read_file_segment(filename, lineno, context)
+    if seg is None:
+        # fallback with single line
+        seg = {
+            "filename": filename,
+            "lineno": lineno,
+            "start": lineno,
+            "end": lineno,
+            "highlight": lineno,
+            "lines": [{"no": lineno, "text": line_of_code or ""}],
+        }
+    return jsonify(seg)
+
+
+@app.route("/api/getMatmulC", methods=["POST"])
+def get_matmul_c():
+    global raw_tensor_data
+    data = request.json or {}
+    uuid = data.get("uuid")
+    if not uuid or uuid not in raw_tensor_data:
+        return jsonify({"error": "Operation not found"}), 404
+    op = raw_tensor_data[uuid]
+    a = op.get("input_data")
+    b = op.get("other_data")
+    if a is None or b is None:
+        return jsonify({"error": "MatMul tensors not available"}), 200
+    try:
+        # Compute C on CPU
+        c = a @ b
+        c_cpu = c.cpu()
+        cmin = float(c_cpu.min().item())
+        cmax = float(c_cpu.max().item())
+        return jsonify(
+            {
+                "shape": list(c_cpu.shape),
+                "min": cmin,
+                "max": cmax,
+                "values": c_cpu.numpy().tolist(),
+            }
+        )
+    except Exception as e:
+        return jsonify({"error": f"MatMul compute failed: {e}"}), 200
+
+
 @app.route("/api/getValue", methods=["POST"])
 def get_value():
     global raw_tensor_data, precomputed_c_values, current_fullscreen_op

From cb663239615f456efe93c7c78ee054843d5adebb Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Fri, 3 Oct 2025 16:27:11 -0400
Subject: [PATCH 34/90] add some ops

---
 triton_viz/core/nki.py | 228 ++++++++++++++++++++++++++++-------------
 1 file changed, 159 insertions(+), 69 deletions(-)

diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index 2a2efaa9..0088001f 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -89,14 +89,15 @@ def __init__(self, buffer=None, name="", **kwargs):
         self.buffer = buffer
         self.name = name
         self.kwargs = kwargs
+        value = None
         if "shape" in kwargs and "dtype" in kwargs:
             shape = kwargs.pop("shape")
             dtype = kwargs.pop("dtype")
-            self._value = np.ndarray(shape, dtype=dtype)
-        elif "value" in kwargs:
-            self._value = kwargs["value"]
-        else:
-            self._value = None
+            value = np.ndarray(shape, dtype=dtype)
+        if "value" in kwargs:
+            assert value is None or value.shape == kwargs["value"].shape
+            value = kwargs["value"]
+        self._value = value
 
     @property
     def shape(self):
@@ -109,11 +110,18 @@ def dtype(self):
     @property
     def value(self):
         return self._value
+    
+    @property
+    def data(self):
+        return self._value
 
     @value.setter
     def value(self, new_value):
         self._value = new_value
 
+    def data_ptr(self):
+        return self.ctypes.data
+
     def __repr__(self):
         return f"NDArray(shape={self.shape}, dtype={self.dtype}, name={self.name})"
 
@@ -125,75 +133,82 @@ def __getitem__(self, keys):
         # Apply the slicing to the underlying numpy array
         new_keys = []
         if isinstance(keys, tuple):
+            arr_dim = 0
             for k in keys:
                 if isinstance(k, NDArray):
-                    new_keys.append(k._value)
+                    dim_len = self._value.shape[arr_dim]
+                    new_keys.append(k._value.clip(0, dim_len - 1))
                 elif isinstance(k, NLSlice):
                     new_keys.append(slice(k.start, k.stop, k.step))
+                elif k is None:
+                    new_keys.append(k)
+                    arr_dim -= 1 # add new dim -> revisit arr_dim for next key
                 else:
                     new_keys.append(k)
+                arr_dim += 1
+        
         sliced_value = self._value[tuple(new_keys)]
 
         # Create a new NDArray with the sliced data
         return NDArray(value=sliced_value, name=f"{self.name}_slice")
 
 
-    # Define operator +/-/*//
-    def __add__(self, other):
+    def _binary_op(self, other, op_func, op_name, op_symbol):
         if isinstance(other, NDArray):
-            return NDArray(value=self._value + other._value, name=f"{self.name}_add_{other.name}")
+            return NDArray(value=op_func(self._value, other._value), name=f"{self.name}_{op_name}_{other.name}")
         elif np.isscalar(other):
-            return NDArray(value=self._value + other, name=f"{self.name}_add_scalar")
-        raise TypeError(f"Unsupported operand type(s) for +: 'NDArray' and '{type(other).__name__}'")
+            return NDArray(value=op_func(self._value, other), name=f"{self.name}_{op_name}_scalar")
+        raise TypeError(f"Unsupported operand type(s) for {op_symbol}: 'NDArray' and '{type(other).__name__}'")
 
-    def __radd__(self, other):
+    def _rbinary_op(self, other, op_func, op_name, op_symbol):
         if isinstance(other, NDArray):
-            return NDArray(value=other._value + self._value, name=f"{other.name}_add_{self.name}")
+            return NDArray(value=op_func(other._value, self._value), name=f"{other.name}_{op_name}_{self.name}")
         elif np.isscalar(other):
-            return NDArray(value=other + self._value, name=f"scalar_add_{self.name}")
-        raise TypeError(f"Unsupported operand type(s) for +: '{type(other).__name__}' and 'NDArray'")
+            return NDArray(value=op_func(other, self._value), name=f"scalar_{op_name}_{self.name}")
+        raise TypeError(f"Unsupported operand type(s) for {op_symbol}: '{type(other).__name__}' and 'NDArray'")
+
+    # Define operator +/-/*//
+    def __add__(self, other):
+        return self._binary_op(other, lambda a, b: a + b, "add", "+")
+
+    def __radd__(self, other):
+        return self._rbinary_op(other, lambda a, b: a + b, "add", "+")
 
     def __sub__(self, other):
-        if isinstance(other, NDArray):
-            return NDArray(value=self._value - other._value, name=f"{self.name}_sub_{other.name}")
-        elif np.isscalar(other):
-            return NDArray(value=self._value - other, name=f"{self.name}_sub_scalar")
-        raise TypeError(f"Unsupported operand type(s) for -: 'NDArray' and '{type(other).__name__}'")
+        return self._binary_op(other, lambda a, b: a - b, "sub", "-")
 
     def __rsub__(self, other):
-        if isinstance(other, NDArray):
-            return NDArray(value=other._value - self._value, name=f"{other.name}_sub_{self.name}")
-        elif np.isscalar(other):
-            return NDArray(value=other - self._value, name=f"scalar_sub_{self.name}")
-        raise TypeError(f"Unsupported operand type(s) for -: '{type(other).__name__}' and 'NDArray'")
+        return self._rbinary_op(other, lambda a, b: a - b, "sub", "-")
 
     def __mul__(self, other):
-        if isinstance(other, NDArray):
-            return NDArray(value=self._value * other._value, name=f"{self.name}_mul_{other.name}")
-        elif np.isscalar(other):
-            return NDArray(value=self._value * other, name=f"{self.name}_mul_scalar")
-        raise TypeError(f"Unsupported operand type(s) for *: 'NDArray' and '{type(other).__name__}'")
+        return self._binary_op(other, lambda a, b: a * b, "mul", "*")
 
     def __rmul__(self, other):
-        if isinstance(other, NDArray):
-            return NDArray(value=other._value * self._value, name=f"{other.name}_mul_{self.name}")
-        elif np.isscalar(other):
-            return NDArray(value=other * self._value, name=f"scalar_mul_{self.name}")
-        raise TypeError(f"Unsupported operand type(s) for *: '{type(other).__name__}' and 'NDArray'")
+        return self._rbinary_op(other, lambda a, b: a * b, "mul", "*")
 
     def __truediv__(self, other):
-        if isinstance(other, NDArray):
-            return NDArray(value=self._value / other._value, name=f"{self.name}_div_{other.name}")
-        elif np.isscalar(other):
-            return NDArray(value=self._value / other, name=f"{self.name}_div_scalar")
-        raise TypeError(f"Unsupported operand type(s) for /: 'NDArray' and '{type(other).__name__}'")
+        return self._binary_op(other, lambda a, b: a / b, "div", "/")
 
     def __rtruediv__(self, other):
-        if isinstance(other, NDArray):
-            return NDArray(value=other._value / self._value, name=f"{other.name}_div_{self.name}")
-        elif np.isscalar(other):
-            return NDArray(value=other / self._value, name=f"scalar_div_{self.name}")
-        raise TypeError(f"Unsupported operand type(s) for /: '{type(other).__name__}' and 'NDArray'")
+        return self._rbinary_op(other, lambda a, b: a / b, "div", "/")
+
+    def __lt__(self, other):
+        return self._binary_op(other, lambda a, b: a < b, "lt", "<")
+
+    def __gt__(self, other):
+        return self._binary_op(other, lambda a, b: a > b, "gt", ">")
+
+    def __le__(self, other):
+        return self._binary_op(other, lambda a, b: a <= b, "le", "<=")
+
+    def __ge__(self, other):
+        return self._binary_op(other, lambda a, b: a >= b, "ge", ">=")
+
+    def __and__(self, other):
+        return self._binary_op(other, lambda a, b: a & b, "and", "&")
+
+    def __or__(self, other):
+        return self._binary_op(other, lambda a, b: a | b, "or", "|")
 
 
 class Builder:
@@ -206,7 +221,7 @@ def __init__(self, grid_dims=None):
         self.fn = None
         self.shared_hbm_arrays = {}
 
-    def set_grid_dims(self, grid_dims):
+    def set_grid_dim(self, grid_dims):
         self.grid_dims = grid_dims
 
     def set_grid_idx(self, x, y, z):
@@ -214,7 +229,7 @@ def set_grid_idx(self, x, y, z):
         self.grid_y = y
         self.grid_z = z
 
-    def ndarray(self, shape, dtype, *, buffer=None, name="", **kwargs):
+    def ndarray(self, shape, dtype, *, buffer=None, name=None, **kwargs):
         if buffer == nl.shared_hbm:
             if name is None:
                 # file name + function name + line number
@@ -234,8 +249,8 @@ def ndarray(self, shape, dtype, *, buffer=None, name="", **kwargs):
             ret = NDArray(buffer=buffer, name=name, shape=shape, dtype=dtype, **kwargs)
         return ret
     
-    def zeros(self, shape, dtype, *, buffer=None, name="", **kwargs):
-        val = np.zeros(shape, dtype=dtype)
+    def zeros(self, shape, dtype, *, buffer=None, name=None, **kwargs):
+        value = np.zeros(shape, dtype=dtype)
         return self.ndarray(
             shape, dtype, buffer=buffer, name=name, value=value, **kwargs
         )
@@ -263,10 +278,21 @@ def program_id(self, axis: int):
 
     def load(self, src: NDArray, *, mask=None, dtype=None, **kwargs):
         value = src._value
-        if mask is not None:
+        if isinstance(mask, NDArray):
+            value = value[mask._value]
+        elif mask is not None:
             value = value[mask]
         if dtype is not None:
             value = value.astype(dtype)
+
+        mask_value = getattr(mask, "_value", np.ones_like(src))
+        new_shape = []
+        for i, _v in enumerate(mask_value.shape):
+            assert np.unique(mask_value.sum(i)).size <= 2
+            new_dim = mask_value.sum(i).flatten()[0]
+            new_shape.append(new_dim)
+
+        value = value.reshape(new_shape)
         return NDArray(value=value, name=src.name, **kwargs)
 
     def load_transpose2d(self, src: NDArray, *, mask=None, dtype=None, **kwargs):
@@ -275,47 +301,111 @@ def load_transpose2d(self, src: NDArray, *, mask=None, dtype=None, **kwargs):
         return self.load(NDArray(value=value.T, name=src.name), mask=mask, dtype=dtype, **kwargs)
 
     def store(self, dst: NDArray, value: NDArray, *, mask=None, **kwargs):
-        if mask is not None:
-            value = value[mask]
-        import tp 
-
-        tp.log(f'{dst=}')
-        tp.log(f'{value=}')
-        dst._value[:] = value._value[:]
+        dst._value[mask._value] = value._value.ravel()
         return dst
 
+    def unary_op(self, x: NDArray, np_func, op_name, **kwargs):
+        return NDArray(value=np_func(x._value), name=f"{x.name}_{op_name}", **kwargs)
+
+    # Elementwise operator implementations
+    def exp(self, x: NDArray, **kwargs):
+        return self._unary_op(x, np.exp, "exp", **kwargs)
+
+    def relu(self, x: NDArray, **kwargs):
+        return self._unary_op(x, lambda v: np.maximum(v, 0), "relu", **kwargs)
+
+    def sigmoid(self, x: NDArray, **kwargs):
+        return self._unary_op(x, lambda v: 1 / (1 + np.exp(-v)), "sigmoid", **kwargs)
+
+    def tanh(self, x: NDArray, **kwargs):
+        return self._unary_op(x, np.tanh, "tanh", **kwargs)
+
+    def silu(self, x: NDArray, **kwargs):
+        # SiLU(x) = x * sigmoid(x)
+        sigmoid_x = 1 / (1 + np.exp(-x._value))
+        return NDArray(value=x._value * sigmoid_x, name=f"{x.name}_silu", **kwargs)
+
+    def gelu(self, x: NDArray, **kwargs):
+        # GELU(x) = 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
+        sqrt_2_pi = np.sqrt(2 / np.pi)
+        inner = sqrt_2_pi * (x._value + 0.044715 * np.power(x._value, 3))
+        return NDArray(value=0.5 * x._value * (1 + np.tanh(inner)), name=f"{x.name}_gelu", **kwargs)
+
+    def sqrt(self, x: NDArray, **kwargs):
+        return self._unary_op(x, np.sqrt, "sqrt", **kwargs)
+
+    def abs(self, x: NDArray, **kwargs):
+        return self._unary_op(x, np.abs, "abs", **kwargs)
+
+    def log(self, x: NDArray, **kwargs):
+        return self._unary_op(x, np.log, "log", **kwargs)
+
+    def pow(self, x: NDArray, exponent, **kwargs):
+        if isinstance(exponent, NDArray):
+            return NDArray(value=np.power(x._value, exponent._value), name=f"{x.name}_pow_{exponent.name}", **kwargs)
+        elif np.isscalar(exponent):
+            return NDArray(value=np.power(x._value, exponent), name=f"{x.name}_pow_{exponent}", **kwargs)
+        else:
+            raise TypeError(f"Unsupported exponent type: {type(exponent)}")
+
+    def reciprocal(self, x: NDArray, **kwargs):
+        return self._unary_op(x, lambda v: 1 / v, "reciprocal", **kwargs)
+
+    def matmul(self, x: NDArray, y: NDArray, transpose_x=False, mask=None, **kwargs):
+        x_value = x._value 
+        if transpose_x:
+            x_value = x_value.T
+        y_value = y._value 
+        return NDArray(value=(x_value @ y_value), name=f"{x.name}_{y.name}_matmul", **kwargs)
+
+    def copy(self, x: NDArray, **kwargs):
+        return self.unary_op(x, np.copy, "copy", **kwargs)
+
+    def range(self, stop):
+        return range(stop)
 
 nki_builder = Builder()
 
 
 def patch():
-    nl.ndarray = lambda *args, **kwargs: nki_builder.ndarray(*args, **kwargs)
-    nl.program_id = lambda axis: nki_builder.program_id(axis)
-    nl.arange = lambda *args: nki_builder.arange(*args)
-    nl.load = lambda src, **kwargs: nki_builder.load(src, **kwargs)
-    nl.store = lambda dst, value, **kwargs: nki_builder.store(dst, value, **kwargs)
+    nl.ndarray = nki_builder.ndarray
+    nl.program_id = nki_builder.program_id
+    nl.arange = nki_builder.arange
+    nl.load = nki_builder.load
+    nl.store = nki_builder.store
     # see https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/nki/api/nki.language.html
 
     # TODO: implement
     # matmul-specific
     # nl.shared_hbm
     # nl.psum
-    nl.affine_range
+    nl.affine_range = nki_builder.range
     nl.par_dim
-    nl.zeros = lambda *args, **kwags: nki_builder.zeros(*args, **kwargs)
+    nl.zeros = nki_builder.zeros
     nl.mgrid
-    #nisa.nc_matmul
-    nl.copy
+    nl.matmul = nki_builder.matmul
+    nl.copy = nki_builder.copy
 
     # attention-specific
-    nl.load_transpose2d = lambda src, **kwargs: nki_builder.load_transpose2d(src, **kwargs)
+    nl.load_transpose2d = nki_builder.load_transpose2d
     #nisa.affine_select
     #nl.tensor_reduce
     #nisa.activation
     nl.broadcast_to
     #nisa.nc_transpose
 
-
+    # Elementwise operators
+    nl.exp = nki_builder.exp
+    nl.relu = nki_builder.relu
+    nl.sigmoid = nki_builder.sigmoid
+    nl.tanh = nki_builder.tanh
+    nl.silu = nki_builder.silu
+    nl.gelu = nki_builder.gelu
+    nl.sqrt = nki_builder.sqrt
+    nl.abs = nki_builder.abs
+    nl.log = nki_builder.log
+    nl.pow = nki_builder.pow
+    nl.reciprocal = nki_builder.reciprocal
 
 
 def unpatch():
@@ -338,7 +428,7 @@ def run(self, *args, **kwargs):
             grid_dims = (grid_dims[0], grid_dims[1], 1)
         elif len(grid_dims) != 3:
             raise ValueError(f"Grid must be 1, 2, or 3 dimensions, got {len(grid_dims)}")
-        nki_builder.set_grid_dims(grid_dims)
+        nki_builder.set_grid_dim(grid_dims)
         nki_builder.shared_hbm_arrays = {}
         nki_builder.fn = self.fn
 

From 7bd6bea5e42f4b3c99e1b5802856b75f4412a528 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Fri, 3 Oct 2025 16:29:57 -0400
Subject: [PATCH 35/90] nki kernels to demo interpreter

---
 nki-examples/matmul.py       | 114 +++++++++++++++++++++++++++
 nki-examples/rmsnorm.py      |  94 ++++++++++++++++++++++
 nki-examples/softmax.py      |  64 ++++++++-------
 nki-examples/test_kernels.py | 146 ++++++++++++++++++++++++-----------
 4 files changed, 347 insertions(+), 71 deletions(-)
 create mode 100644 nki-examples/matmul.py
 create mode 100644 nki-examples/rmsnorm.py

diff --git a/nki-examples/matmul.py b/nki-examples/matmul.py
new file mode 100644
index 00000000..e0eab83d
--- /dev/null
+++ b/nki-examples/matmul.py
@@ -0,0 +1,114 @@
+from neuronxcc import nki
+import neuronxcc.nki.language as nl
+
+import tp
+import torch
+import triton_viz
+import torch_xla.core.xla_model as xm
+from triton_viz.clients import Tracer
+from triton_viz.core import config as cfg
+from triton_viz.core.trace import launches
+import numpy as np
+import math
+
+def matmul_kernel(lhs, rhs):
+    """NKI matmul_kernel to compute a matrix multiplication operation in a tiled manner
+
+    Args:
+        lhs: an input tensor of shape [K,M], where both K and M are multiples for
+          128.  It is the left-hand-side argument of the matrix multiplication,
+          delivered transposed for optimal performance.
+        rhs: an input tensor of shape [K,N], where K is a multiple of 128, and N
+          is a multiple of 512.  It is the right-hand-side argument of the
+          matrix multiplication.
+    Returns:
+        result: the resulting output tensor of shape [M,N]
+    """
+
+    M, K = lhs.shape
+    K_, N = rhs.shape
+    assert K == K_, "lhs and rhs must have the same contraction dimension"
+    result = nl.ndarray((M, N), dtype=lhs.dtype, buffer=nl.shared_hbm)
+
+    TILE_M = 2
+    TILE_K = 2
+    TILE_N = 4
+
+    # Use affine_range to loop over tiles
+    m_range = nl.arange(M)[:, None]
+    n_range = nl.arange(N)[None, :]
+    for m in nl.affine_range(M // TILE_M):
+        for n in nl.affine_range(N // TILE_N):
+            # Allocate a tensor in PSUM
+            res_psum = nl.zeros((TILE_M, TILE_N), nl.int32, buffer=nl.psum)
+
+            for k in nl.affine_range(K // TILE_K):
+                # Declare the tiles on SBUF
+                lhs_tile = nl.ndarray((TILE_K, TILE_M), dtype=lhs.dtype, buffer=nl.sbuf)
+                rhs_tile = nl.ndarray((TILE_K, TILE_N), dtype=rhs.dtype, buffer=nl.sbuf)
+
+                # Load tiles from lhs and rhs
+                lhs_p = nl.arange(TILE_M)[:, None] + m * TILE_M
+                lhs_f = nl.arange(TILE_K)[None, :] + k * TILE_K
+                lhs_tile = nl.load(lhs[lhs_p, lhs_f], mask=(lhs_p < M) & (lhs_f < K))
+                rhs_p = nl.arange(TILE_K)[:, None] + k * TILE_K
+                rhs_f = nl.arange(TILE_N)[None, :] + n * TILE_N
+                rhs_tile = nl.load(rhs[rhs_p, rhs_f], mask=(rhs_p < K) & (rhs_f < N))
+
+                # Accumulate partial-sums into PSUM
+                res_psum += nl.matmul(lhs_tile[...], rhs_tile[...], transpose_x=False)
+
+            # Copy the result from PSUM back to SBUF, and cast to expected output data-type
+            out_mask = (m * TILE_M <= m_range) & (m_range < (m + 1) * TILE_M)
+            out_mask &= (n * TILE_N <= n_range) & (n_range < (n + 1) * TILE_N)
+            res_sb = nl.copy(res_psum, dtype=result.dtype)
+
+            #nl.store(result[m * TILE_M:(m + 1) * TILE_M, n * TILE_N:(n + 1) * TILE_N],
+            #         value=res_sb) # TODO: uncomment this
+            nl.store(result, value=res_sb, mask=out_mask) # TODO: remove this - incorrect NKI syntax but currently needed for interpreter
+
+    return result
+
+TRITON_VIZ = True
+kernel_grid = (1, 1, 1)
+lhs_small = np.arange(64).astype(np.float32).reshape(8, 8)
+rhs_small = np.arange(128).astype(np.float32).reshape(8, 16)
+#lhs_small = np.arange(16).astype(np.float32).reshape(4, 4)
+#rhs_small = np.arange(32).astype(np.float32).reshape(4, 8)
+kernel_args = (lhs_small, rhs_small)
+
+if TRITON_VIZ:
+    tp.log('Executing matmul_kernel with NKI interpreter...')
+    matmul_kernel = triton_viz.trace(clients=Tracer(), backend="nki")(matmul_kernel)
+    kk = matmul_kernel[kernel_grid]
+    z2 = kk(*kernel_args)
+
+    #print(f"Number of launches: {len(launches)}")
+    #if launches:
+    #    launch = launches[-1]
+    #    print(f"Number of records: {len(launch.records)}")
+    #    for i, record in enumerate(launch.records):
+    #        print(f"Record {i}: {type(record).__name__}")
+    #        if hasattr(record, "ptr"):
+    #            print(f"  ptr: {record.ptr}")
+    #        if hasattr(record, "offsets"):
+    #            print(f"  offsets shape: {record.offsets.shape}")
+    #        if hasattr(record, "masks"):
+    #            print(f"  masks shape: {record.masks.shape}")
+
+    ## Try to launch visualization
+    #try:
+    #    triton_viz.launch(share=False)
+    #except Exception as e:
+    #    print(f"\nError during visualization: {e}")
+    #    import traceback
+
+    #    traceback.print_exc()
+else:
+    tp.log('Executing NKI JIT-ed matmul_kernel...')
+    matmul_kernel = nki.jit(matmul_kernel)
+    z2 = nki.simulate_kernel(matmul_kernel[kernel_grid], *kernel_args)
+
+z1 = lhs_small @ rhs_small
+print(np.max(np.abs(z1-z2)))
+assert np.allclose(z1, z2)
diff --git a/nki-examples/rmsnorm.py b/nki-examples/rmsnorm.py
new file mode 100644
index 00000000..854d6304
--- /dev/null
+++ b/nki-examples/rmsnorm.py
@@ -0,0 +1,94 @@
+from torch_xla.core import xla_model as xm
+import neuronxcc.nki as nki
+import neuronxcc.nki.language as nl
+import numpy as np
+import math
+import torch
+
+@nki.jit
+def nki_rmsnorm_kernel(a_tensor, g_tensor):
+  # Calculate out_tensor = a_tensor/RMS(a_tensor) * g_tensor
+  # Where RMS(a_tensor) = sqrt((1/N) * sum(a_tensor * a_tensor))
+  # and N = a_tensor.shape[1]
+  # Reduction (mean) is performed in the free (2nd) dimension
+  B, D = a_tensor.shape
+  out_tensor = nl.ndarray(a_tensor.shape, dtype=a_tensor.dtype,
+                          buffer=nl.shared_hbm)
+
+  # Make sure shapes match
+  assert D == g_tensor.shape[0]
+
+  # Generate tensor indices to index input tensor
+  ix = nl.arange(2)[:, None]
+  iw = nl.arange(1)[:, None]
+  #iy = nl.arange(a_tensor.shape[1])[None, :]
+  iy = nl.arange(8)[None, :]
+
+  # Load RMSNorm weight once, reused by rows/tiles of a_tensor
+  g_tile = nl.load(g_tensor.reshape((1, D))[iw, iy], mask=iy < D)
+
+  # Process 2 rows at a time due to 2-partition tile size limitation
+  # Since we're not reducing across the first dimension
+  # Tiles can be processed independently
+  for i in nl.affine_range(math.ceil(B/2)):
+
+    # Load input data from external memory to on-chip memory
+    mask = (i * 2 + ix < B) & (iy < D)
+    a_tile = nl.load(a_tensor[i * 2 + ix, iy],
+                    mask=mask)
+
+    # Compute element-wise square of a_tensor
+    in_square = nl.square(a_tile)
+
+    # Calculate sum of squared elements, along last dimension
+    square_sum = nl.sum(in_square, axis=[1], mask=mask)
+
+    # Scale and get a reciprocal
+    mean = square_sum / D
+
+    # Take square root of mean and then reciprocal with
+    # rsqrt API (one ISA instruction)
+    rms_reciprocal = nl.rsqrt(mean)
+
+    # Scale the input tensor
+    out_tile = nl.multiply(a_tile, rms_reciprocal)
+
+    # Broadcast weight along first axis to match tensor shape
+    # B_active = min(B - i * 2, 2)
+    g_bcast = g_tile.broadcast_to((2, 8))
+
+    # Multiply with the RMSNorm weight
+    out_tile[...] = nl.multiply(out_tile, g_bcast,
+                           mask=(i * 2 + ix < B))
+
+    # store the addition results back to external memory (out_tensor)
+    nl.store(out_tensor[i * 2 + ix, iy], value=out_tile,
+            mask=mask)
+
+  return out_tensor
+
+# ref
+def torch_rmsnorm_kernel(a_tensor, g_tensor):
+  # Square the tensor (element-wise)
+  in_square = a_tensor.pow(2)
+  # Calculate means in the free dimension
+  mean = in_square.mean(dim=1, keepdim=True)
+  # Scale by reciprocal of sqrt(mean)
+  tensor = a_tensor * torch.rsqrt(mean)
+
+  # Scale the output by the weight
+  return tensor * g_tensor
+
+device = 'cpu'
+
+a_tensor = torch.arange(15).float().view(3, 5).to(device=device)
+g_tensor = torch.arange(5).float().to(device=device)
+
+output_nki = nki.simulate_kernel(nki_rmsnorm_kernel, a_tensor.numpy(), g_tensor.numpy())
+
+output_torch = torch_rmsnorm_kernel(a_tensor, g_tensor).numpy()
+
+if np.allclose(output_torch, output_nki, atol=1e-5, rtol=1e-3):
+  print("NKI and Torch match")
+else:
+  print("NKI and Torch differ")
diff --git a/nki-examples/softmax.py b/nki-examples/softmax.py
index f84b2711..628ee4b6 100644
--- a/nki-examples/softmax.py
+++ b/nki-examples/softmax.py
@@ -7,9 +7,10 @@
 from triton_viz.clients import Tracer
 from triton_viz.core import config as cfg
 from triton_viz.core.trace import launches
+import numpy as np
+import math
 
 
-@triton_viz.trace(clients=Tracer())
 @nki.jit
 def softmax(in_tensor):
     # assume B,D; softmax on D for now
@@ -33,31 +34,38 @@ def softmax(in_tensor):
 if __name__ == "__main__":
     cfg.reset()
     device = "cpu"
-    size = 16
-    BLOCK_SIZE = 8
+    B, D = 16, 32
     torch.manual_seed(0)
-    x = torch.arange(size, dtype=torch.float32, device=device)
-    softmax[(1,1,1)](x)
-
-    # Print records to see what's being captured
-    print(f"Number of launches: {len(launches)}")
-    if launches:
-        launch = launches[-1]
-        print(f"Number of records: {len(launch.records)}")
-        for i, record in enumerate(launch.records):
-            print(f"Record {i}: {type(record).__name__}")
-            if hasattr(record, "ptr"):
-                print(f"  ptr: {record.ptr}")
-            if hasattr(record, "offsets"):
-                print(f"  offsets shape: {record.offsets.shape}")
-            if hasattr(record, "masks"):
-                print(f"  masks shape: {record.masks.shape}")
-
-    # Try to launch visualization
-    try:
-        triton_viz.launch()
-    except Exception as e:
-        print(f"\nError during visualization: {e}")
-        import traceback
-
-        traceback.print_exc()
+    x = torch.rand((B, D), dtype=torch.float32, device=device)
+    y = np.exp(x) / np.exp(x).sum(-1, keepdim=True)
+    TRITON_VIZ = False
+
+    if TRITON_VIZ:
+        softmax = triton_viz.trace()(softmax)
+        softmax[(1,1,1)](x)
+
+        # Print records to see what's being captured
+        print(f"Number of launches: {len(launches)}")
+        if launches:
+            launch = launches[-1]
+            print(f"Number of records: {len(launch.records)}")
+            for i, record in enumerate(launch.records):
+                print(f"Record {i}: {type(record).__name__}")
+                if hasattr(record, "ptr"):
+                    print(f"  ptr: {record.ptr}")
+                if hasattr(record, "offsets"):
+                    print(f"  offsets shape: {record.offsets.shape}")
+                if hasattr(record, "masks"):
+                    print(f"  masks shape: {record.masks.shape}")
+
+        # Try to launch visualization
+        try:
+            triton_viz.launch()
+        except Exception as e:
+            print(f"\nError during visualization: {e}")
+            import traceback
+
+            traceback.print_exc()
+    else:
+        y_pred = nki.simulate_kernel(softmax, x.numpy())
+        assert np.allclose(y, y_pred)
diff --git a/nki-examples/test_kernels.py b/nki-examples/test_kernels.py
index d4426918..a3e05657 100644
--- a/nki-examples/test_kernels.py
+++ b/nki-examples/test_kernels.py
@@ -1,14 +1,15 @@
 import math
+import tp
 import neuronxcc.nki.language as nl
 import neuronxcc.nki as nki
 import numpy as np
 import torch
 from triton_viz.clients import Tracer
 import triton_viz
+from triton_viz.core.trace import launches
 
-TRITON_VIZ = True
 
-def add_kernel(a, b): # fails @ ix < B
+def add_kernel(a, b):  # fails @ ix < B
     pid_x = nl.program_id(0)
     pid_y = nl.program_id(1)
     B, D = a.shape
@@ -25,7 +26,7 @@ def add_kernel(a, b): # fails @ ix < B
     nl.store(c_output[ix, iy], value=c_tmp, mask=mask)
     return c_output
 
-def copy_kernel(a): # fails @ ix < B
+def copy_kernel(a):  # fails @ ix < B
     B, D = a.shape
     out = nl.ndarray(a.shape, dtype=a.dtype, buffer=nl.shared_hbm)
 
@@ -34,42 +35,53 @@ def copy_kernel(a): # fails @ ix < B
 
     ix = pid_x * 128 + nl.arange(128)[:, None]
     iy = pid_y * 512 + nl.arange(512)[None, :]
-    #mask = (ix < B) & (iy < D)
+    mask = (ix < B) & (iy < D)
 
-    #a_tmp = nl.load(a[ix, iy], mask=mask)
-    #nl.store(out[ix, iy], value=a_tmp, mask=mask)
-    a_tmp = nl.load(a[ix, iy])
-    nl.store(out[ix, iy], value=a_tmp)
+    a_tmp = nl.load(a[ix, iy], mask=mask)
+    nl.store(out[ix, iy], value=a_tmp, mask=mask)
     return out
 
-def print_kernel(): # works
+def print_kernel():  # works
     a = nl.ndarray([4, 4], dtype=nl.float32, buffer=nl.shared_hbm)
-    y = nl.ndarray([4, 4], dtype=np.float32,)
-    nl.store(a, value=y) 
+    y = nl.ndarray(
+        [4, 4],
+        dtype=np.float32,
+    )
+    nl.store(a, value=y)
     print(a)
     return a
 
-def tmp0_kernel(a): # works
+def tmp0_kernel(a):  # works
     B, D = a.shape
     out = nl.ndarray(a.shape, dtype=a.dtype, buffer=nl.shared_hbm)
 
     pid_x = nl.program_id(0)
     pid_y = nl.program_id(1)
-    import tp
-    tp.log(f'{pid_x=}, {pid_y=}')
+    nl.device_print(f"pid_x:", pid_x);
+    nl.device_print(f"pid_y:", pid_y);
 
-    ix = pid_x * 128 + nl.arange(128)[:, None]
-    iy = pid_y * 512 + nl.arange(512)[None, :]
-    #mask = (ix < B) & (iy < D)
+    ix = pid_x * 2 + nl.arange(2)[:, None]
+    iy = pid_y * 4 + nl.arange(4)[None, :]
+    mask = (ix < B) & (iy < D)
+    a_tmp = nl.load(a[ix, iy], mask=mask)
+    nl.device_print(f"a_tmp:", a_tmp)
+
+    ix2 = pid_x * 3 + nl.arange(3)[:, None]
+    iy2 = pid_y * 4 + nl.arange(4)[None, :]
+    mask2 = (ix2 < B) & (iy2 < D)
+    a_tmp2 = nl.load(a[ix2, iy2], mask=mask2)
+    nl.device_print(f"a_tmp2:", a_tmp2)
+
+    iy3 = nl.arange(3)[None, :] < -1
+    a_tmp3 = nl.load(a[80, nl.arange(3)[None, :]], mask=iy3)
+    nl.device_print(f"a_tmp3:", a_tmp3)
 
-    #a_tmp = nl.load(a[ix, iy], mask=mask)
-    #nl.store(out[ix, iy], value=a_tmp, mask=mask)
-    #a_tmp = nl.load(a[:128, :128])
-    a_tmp = nl.load(a[ix, iy])
-    nl.store(out[ix, iy], value=a_tmp)
+    nl.store(out[ix, iy], value=a_tmp, mask=mask)
+    # for load(src, mask), src.shape, mask.shape need to be same, return shape
+    # for store(dst, value, mask), dst.shape, value.shape, mask.shape need to be same, return shape
     return out
 
-def xyz_kernel(a): # works
+def xyz_kernel(a):  # works
     B, T, C, H, W = a.shape
     out = nl.ndarray(a.shape, dtype=a.dtype, buffer=nl.shared_hbm)
 
@@ -77,36 +89,37 @@ def xyz_kernel(a): # works
     pid_y = nl.program_id(1)
     pid_z = nl.program_id(2)
 
-
     for i in range(H // 128):
         for j in range(W // 512):
-            #i_h = pid_x * 128 + nl.arange(128)[None, None, None, :, None]
-            #i_n = pid_y * 512 + nl.arange(512)[None, None, None, None, :]
             i_h = pid_x * 128 + nl.arange(128)[:, None]
             i_n = pid_y * 512 + nl.arange(512)[None, :]
-            a_tmp = nl.load(a[pid_x, pid_y, pid_z, i_h, i_n])
-            nl.store(out[pid_x, pid_y, pid_z, i_h, i_n], value=a_tmp)
+            mask = (i_h < H) & (i_n < W)
+            a_tmp = nl.load(a[pid_x, pid_y, pid_z, i_h, i_n], mask=mask)
+            nl.store(out[pid_x, pid_y, pid_z, i_h, i_n], value=a_tmp, mask=mask)
     return out
 
-def xyz_kernel(a): # works
-    B, T, C, H, W = a.shape
+def tmp1_kernel(a):
+    B, D = a.shape
     out = nl.ndarray(a.shape, dtype=a.dtype, buffer=nl.shared_hbm)
 
     pid_x = nl.program_id(0)
     pid_y = nl.program_id(1)
-    pid_z = nl.program_id(2)
 
+    ix = pid_x * 2 + nl.arange(2)[:, None]
+    iy = pid_y * 4 + nl.arange(4)[None, :]
+    mask = (ix < B) & (iy < D)
+    a_tmp = nl.load(a[ix, iy], mask=mask)
 
-    for i in range(H // 128):
-        for j in range(W // 512):
-            #i_h = pid_x * 128 + nl.arange(128)[None, None, None, :, None]
-            #i_n = pid_y * 512 + nl.arange(512)[None, None, None, None, :]
-            i_h = pid_x * 128 + nl.arange(128)[:, None]
-            i_n = pid_y * 512 + nl.arange(512)[None, :]
-            a_tmp = nl.load(a[pid_x, pid_y, pid_z, i_h, i_n])
-            nl.store(out[pid_x, pid_y, pid_z, i_h, i_n], value=a_tmp)
+    #iy3 = nl.arange(3)[None, :] < -1
+    #iy3 = nl.mgrid[80:83, 80:85] < -1
+    iy3 = (nl.arange(3)[:, None] < -1) & (nl.arange(5)[None, :] < -1)
+    a_tmp3 = nl.load(a[80:83, 80:85], mask=iy3)
+    nl.device_print(f"a_tmp3:", a_tmp3)
+
+    nl.store(out[ix, iy], value=a_tmp, mask=mask)
     return out
 
+
 B, D = 1024, 1024
 x = torch.rand((B, D))
 y = torch.rand((B, D))
@@ -114,7 +127,7 @@ def xyz_kernel(a): # works
 blocks_x = math.ceil(B / 128)
 blocks_y = math.ceil(D / 512)
 
-kernel = tmp0_kernel
+kernel = tmp1_kernel
 if kernel == add_kernel:
     kernel_grid = (blocks_x, blocks_y)
     kernel_args = (x.numpy(), y.numpy())
@@ -124,10 +137,21 @@ def xyz_kernel(a): # works
     kernel_args = (x.numpy(),)
     z1 = x
 elif kernel == print_kernel:
-    kernel_grid = (1,1,1)
+    kernel_grid = (1, 1, 1)
     kernel_args = ()
     z1 = x
 if kernel == tmp0_kernel:
+    #B, D = 129, 512
+    B, D = 3, 5
+    #x = torch.rand((B, D))
+    #y = torch.rand((B, D))
+    x = torch.arange(B*D).int().reshape(B, D)
+    y = -torch.arange(B*D).int().reshape(B, D)
+
+    #blocks_x = math.ceil(B / 128)
+    blocks_x = math.ceil(B / 2)
+    blocks_y = math.ceil(D / 4)
+
     kernel_grid = (blocks_x, blocks_y)
     kernel_args = (x.numpy(),)
     z1 = x
@@ -137,12 +161,48 @@ def xyz_kernel(a): # works
     kernel_grid = (B, T, C)
     kernel_args = (x.numpy(),)
     z1 = x
+if kernel == tmp1_kernel:
+    B, D = 3, 5
+    x = torch.arange(B*D).int().reshape(B, D)
+    y = -torch.arange(B*D).int().reshape(B, D)
+
+    blocks_x = math.ceil(B / 2)
+    blocks_y = math.ceil(D / 4)
+
+    kernel_grid = (blocks_x, blocks_y)
+    kernel_args = (x.numpy(),)
+    z1 = x
+
+TRITON_VIZ = False
 
 if TRITON_VIZ:
-    kernel = triton_viz.trace(clients=Tracer(), backend='nki')(kernel)
-    z2 = kernel[kernel_grid](*kernel_args)
+    kernel = triton_viz.trace(clients=Tracer(), backend="nki")(kernel)
+    kk = kernel[kernel_grid]
+    z2 = kk(*kernel_args)
     z2 = torch.from_numpy(z2)
     print((z1 - z2).abs().max())
+
+    print(f"Number of launches: {len(launches)}")
+    if launches:
+        launch = launches[-1]
+        print(f"Number of records: {len(launch.records)}")
+        for i, record in enumerate(launch.records):
+            print(f"Record {i}: {type(record).__name__}")
+            if hasattr(record, "ptr"):
+                print(f"  ptr: {record.ptr}")
+            if hasattr(record, "offsets"):
+                print(f"  offsets shape: {record.offsets.shape}")
+            if hasattr(record, "masks"):
+                print(f"  masks shape: {record.masks.shape}")
+
+    # Try to launch visualization
+    try:
+        triton_viz.launch(share=False)
+    except Exception as e:
+        print(f"\nError during visualization: {e}")
+        import traceback
+
+        traceback.print_exc()
 else:
     kernel = nki.jit(kernel)
     z2 = nki.simulate_kernel(kernel[kernel_grid], *kernel_args)

From 9afefc6c623d5741e922b3d4dbe77320d71a35e4 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Fri, 3 Oct 2025 16:30:27 -0400
Subject: [PATCH 36/90] fix misc tests

---
 tests/test_core.py       | 4 ++--
 triton_viz/core/trace.py | 3 +--
 2 files changed, 3 insertions(+), 4 deletions(-)

diff --git a/tests/test_core.py b/tests/test_core.py
index 5ae38bb1..d8b9482f 100644
--- a/tests/test_core.py
+++ b/tests/test_core.py
@@ -47,9 +47,9 @@ def my_kernel(x_ptr, y_ptr, out_ptr, BLOCK_SIZE: tl.constexpr):
         tl.store(out_ptr + offs, tl.load(x_ptr + offs) + tl.load(y_ptr + offs))
 
     # Should be wrapped as a Trace object.
-    from triton_viz.core.trace import Trace
+    from triton_viz.core.trace import TritonTrace
 
-    assert isinstance(my_kernel, Trace)
+    assert isinstance(my_kernel, TritonTrace)
 
     # Verify client de-duplication and addition logic
     clients = my_kernel.client_manager.clients
diff --git a/triton_viz/core/trace.py b/triton_viz/core/trace.py
index 61745c0b..344642f3 100644
--- a/triton_viz/core/trace.py
+++ b/triton_viz/core/trace.py
@@ -94,8 +94,7 @@ def __call__(self, *args, **kwargs):
         return self[(1, 1, 1)](*args, **kwargs)
 
 
-#def trace(clients: str | Client | None = None, backend: str = "triton"):
-def trace(clients: str | Client | None = None, backend: str = "nki"):
+def trace(clients: str | Client | None = None, backend: str = "triton"):
     """
     Create a trace object that can be used to run a kernel with instrumentation clients.
 

From 95ef0b69905e432c2d08c7f3bf30719b2e43e1d6 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Fri, 3 Oct 2025 16:31:41 -0400
Subject: [PATCH 37/90] prepare for nki tracing

---
 triton_viz/clients/tracer/tracer.py |  2 +-
 triton_viz/core/patch.py            | 35 ++++++++++++++++++++++++-----
 2 files changed, 30 insertions(+), 7 deletions(-)

diff --git a/triton_viz/clients/tracer/tracer.py b/triton_viz/clients/tracer/tracer.py
index 3747cd16..dd6e3292 100644
--- a/triton_viz/clients/tracer/tracer.py
+++ b/triton_viz/clients/tracer/tracer.py
@@ -58,7 +58,7 @@ def grid_callback(self, grid: tuple[int]):
 
     def register_op_callback(self, op_type: type[Op]) -> OpCallbacks:
         def pre_load_callback(
-            ptr, mask, other, cache_modifier, eviction_policy, is_volatile
+            ptr, mask, *ignore_args, **ignore_kwargs
         ):
             if not self.sample:
                 return
diff --git a/triton_viz/core/patch.py b/triton_viz/core/patch.py
index 8658e1c3..c99a5563 100644
--- a/triton_viz/core/patch.py
+++ b/triton_viz/core/patch.py
@@ -39,6 +39,7 @@
 )
 from triton.runtime.interpreter import _patch_lang as triton_patch_lang
 from triton.runtime import JITFunction
+from triton_viz.core.nki import nki_builder
 
 op_list = [
     ProgramId,
@@ -64,6 +65,13 @@
     Idiv,
     Rsqrt,
     CastImpl,
+
+    #ProgramId,
+    #Store,
+    #Load,
+    #Dot,
+    #UnaryOp,
+    #MakeRange
 ]
 original_ops = {
     ProgramId: interpreter_builder.create_get_program_id,
@@ -86,6 +94,25 @@
     Idiv: interpreter_builder.create_idiv,
     Rsqrt: interpreter_builder.create_rsqrt,
     CastImpl: interpreter_builder.cast_impl,
+
+    #ProgramId: nki_builder.program_id,
+    #Store: nki_builder.store,
+    #Load: nki_builder.load,
+    #Dot: nki_builder.matmul,
+    #UnaryOp: nki_builder.unary_op,
+    ##BinaryOp: nki_builder.binary_op,
+    ##TernaryOp: nki_builder.ternary_op,
+    #MakeRange: nki_builder.arange,
+    ##AddPtr: nki_builder.create_addptr,
+    ##ExpandDims: nki_builder.create_expand_dims,
+    ##Broadcast: nki_builder.create_broadcast,
+    ##Splat: nki_builder.create_splat,
+    ##MakeBlockPointer: nki_builder.create_make_block_ptr,
+    ##TensorPointerLoad: nki_builder.create_tensor_pointer_load,
+    ##TensorPointerStore: nki_builder.create_tensor_pointer_store,
+    ##Idiv: nki_builder.create_idiv,
+    ##Rsqrt: nki_builder.create_rsqrt,
+    ##CastImpl: nki_builder.cast_impl,
 }
 reduce_map: dict[type[Op], Callable] = {
     ReduceMax: tl.max,
@@ -143,16 +170,12 @@ def patch_op(op_type: type[Op], callbacks: OpCallbacks):
         op_name = original_ops[op_type].__name__
         current_op = getattr(interpreter_builder, op_name)
         patched_op = PatchOp(current_op, op_type, callbacks)
-        setattr(
-            interpreter_builder,
-            op_name,
-            lambda *args, **kwargs: patched_op(*args, **kwargs),
-        )
+        setattr(interpreter_builder, op_name, patched_op)
     elif op_type in reduce_map:
         op_name = reduce_map[op_type].__name__
         current_op = getattr(tl, op_name)
         patched_op = PatchOp(current_op, op_type, callbacks)
-        setattr(tl, op_name, lambda *args, **kwargs: patched_op(*args, **kwargs))
+        setattr(tl, op_name, patched_op)
     else:
         raise ValueError(f"Patching operator {op_type} not supported")
 

From 2636088944f97f405185a0ea9925ffc50150ec04 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Fri, 3 Oct 2025 16:31:59 -0400
Subject: [PATCH 38/90] uv

---
 uv.lock | 795 ++++++++++++++++++++++++++++++++++++++++++++++++++++++++
 1 file changed, 795 insertions(+)
 create mode 100644 uv.lock

diff --git a/uv.lock b/uv.lock
new file mode 100644
index 00000000..b2a5d4cb
--- /dev/null
+++ b/uv.lock
@@ -0,0 +1,795 @@
+version = 1
+requires-python = ">=3.7"
+
+[[package]]
+name = "anytree"
+version = "2.8.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "six" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/d8/45/de59861abc8cb66e9e95c02b214be4d52900aa92ce34241a957dcf1d569d/anytree-2.8.0.tar.gz", hash = "sha256:3f0f93f355a91bc3e6245319bf4c1d50e3416cc7a35cc1133c1ff38306bbccab", size = 189484 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/a8/65/be23d8c3ecd68d40541d49812cd94ed0f3ee37eb88669ca15df0e43daed1/anytree-2.8.0-py2.py3-none-any.whl", hash = "sha256:14c55ac77492b11532395049a03b773d14c7e30b22aa012e337b1e983de31521", size = 41717 },
+]
+
+[[package]]
+name = "cairocffi"
+version = "1.6.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "cffi" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/6e/37/982a8650ab8c0b3bd0021e3f7041e9707ce35f6cc43b7def2c43ff109dd9/cairocffi-1.6.1.tar.gz", hash = "sha256:78e6bbe47357640c453d0be929fa49cd05cce2e1286f3d2a1ca9cbda7efdb8b7", size = 87436 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/17/be/a5d2c16317c6a890502725970589ae7f06cfc66b2e6916ba0a86973403c8/cairocffi-1.6.1-py3-none-any.whl", hash = "sha256:aa78ee52b9069d7475eeac457389b6275aa92111895d78fbaa2202a52dac112e", size = 75096 },
+]
+
+[[package]]
+name = "certifi"
+version = "2025.8.3"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/dc/67/960ebe6bf230a96cda2e0abcf73af550ec4f090005363542f0765df162e0/certifi-2025.8.3.tar.gz", hash = "sha256:e564105f78ded564e3ae7c923924435e1daa7463faeab5bb932bc53ffae63407", size = 162386 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/e5/48/1549795ba7742c948d2ad169c1c8cdbae65bc450d6cd753d124b17c8cd32/certifi-2025.8.3-py3-none-any.whl", hash = "sha256:f6c12493cfb1b06ba2ff328595af9350c65d6644968e5d3a2ffd78699af217a5", size = 161216 },
+]
+
+[[package]]
+name = "cffi"
+version = "1.15.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "pycparser" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/2b/a8/050ab4f0c3d4c1b8aaa805f70e26e84d0e27004907c5b8ecc1d31815f92a/cffi-1.15.1.tar.gz", hash = "sha256:d400bfb9a37b1351253cb402671cea7e89bdecc294e8016a707f6d1d8ac934f9", size = 508501 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/e8/ff/c4b7a358526f231efa46a375c959506c87622fb4a2c5726e827c55e6adf2/cffi-1.15.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:39d39875251ca8f612b6f33e6b1195af86d1b3e60086068be9cc053aa4376e21", size = 179233 },
+    { url = "https://files.pythonhosted.org/packages/ea/be/c4ad40ad441ac847b67c7a37284ae3c58f39f3e638c6b0f85fb662233825/cffi-1.15.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:285d29981935eb726a4399badae8f0ffdff4f5050eaa6d0cfc3f64b857b77185", size = 174280 },
+    { url = "https://files.pythonhosted.org/packages/ed/a3/c5f01988ddb70a187c3e6112152e01696188c9f8a4fa4c68aa330adbb179/cffi-1.15.1-cp310-cp310-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:3eb6971dcff08619f8d91607cfc726518b6fa2a9eba42856be181c6d0d9515fd", size = 421712 },
+    { url = "https://files.pythonhosted.org/packages/ef/41/19da352d341963d29a33bdb28433ba94c05672fb16155f794fad3fd907b0/cffi-1.15.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:21157295583fe8943475029ed5abdcf71eb3911894724e360acff1d61c1d54bc", size = 449886 },
+    { url = "https://files.pythonhosted.org/packages/af/da/9441d56d7dd19d07dcc40a2a5031a1f51c82a27cee3705edf53dadcac398/cffi-1.15.1-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:5635bd9cb9731e6d4a1132a498dd34f764034a8ce60cef4f5319c0541159392f", size = 450520 },
+    { url = "https://files.pythonhosted.org/packages/aa/02/ab15b3aa572759df752491d5fa0f74128cd14e002e8e3257c1ab1587810b/cffi-1.15.1-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:2012c72d854c2d03e45d06ae57f40d78e5770d252f195b93f581acf3ba44496e", size = 446015 },
+    { url = "https://files.pythonhosted.org/packages/88/89/c34caf63029fb7628ec2ebd5c88ae0c9bd17db98c812e4065a4d020ca41f/cffi-1.15.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:dd86c085fae2efd48ac91dd7ccffcfc0571387fe1193d33b6394db7ef31fe2a4", size = 441830 },
+    { url = "https://files.pythonhosted.org/packages/32/bd/d0809593f7976828f06a492716fbcbbfb62798bbf60ea1f65200b8d49901/cffi-1.15.1-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:fa6693661a4c91757f4412306191b6dc88c1703f780c8234035eac011922bc01", size = 434743 },
+    { url = "https://files.pythonhosted.org/packages/0e/65/0d7b5dad821ced4dcd43f96a362905a68ce71e6b5f5cfd2fada867840582/cffi-1.15.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:59c0b02d0a6c384d453fece7566d1c7e6b7bae4fc5874ef2ef46d56776d61c9e", size = 464113 },
+    { url = "https://files.pythonhosted.org/packages/9f/52/1e2b43cfdd7d9a39f48bc89fcaee8d8685b1295e205a4f1044909ac14d89/cffi-1.15.1-cp310-cp310-win32.whl", hash = "sha256:cba9d6b9a7d64d4bd46167096fc9d2f835e25d7e4c121fb2ddfc6528fb0413b2", size = 170412 },
+    { url = "https://files.pythonhosted.org/packages/0e/e2/a23af3d81838c577571da4ff01b799b0c2bbde24bd924d97e228febae810/cffi-1.15.1-cp310-cp310-win_amd64.whl", hash = "sha256:ce4bcc037df4fc5e3d184794f27bdaab018943698f4ca31630bc7f84a7b69c6d", size = 179060 },
+    { url = "https://files.pythonhosted.org/packages/23/8b/2e8c2469eaf89f7273ac685164949a7e644cdfe5daf1c036564208c3d26b/cffi-1.15.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:3d08afd128ddaa624a48cf2b859afef385b720bb4b43df214f85616922e6a5ac", size = 179198 },
+    { url = "https://files.pythonhosted.org/packages/f9/96/fc9e118c47b7adc45a0676f413b4a47554e5f3b6c99b8607ec9726466ef1/cffi-1.15.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:3799aecf2e17cf585d977b780ce79ff0dc9b78d799fc694221ce814c2c19db83", size = 174221 },
+    { url = "https://files.pythonhosted.org/packages/10/72/617ee266192223a38b67149c830bd9376b69cf3551e1477abc72ff23ef8e/cffi-1.15.1-cp311-cp311-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:a591fe9e525846e4d154205572a029f653ada1a78b93697f3b5a8f1f2bc055b9", size = 441694 },
+    { url = "https://files.pythonhosted.org/packages/91/bc/b7723c2fe7a22eee71d7edf2102cd43423d5f95ff3932ebaa2f82c7ec8d0/cffi-1.15.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3548db281cd7d2561c9ad9984681c95f7b0e38881201e157833a2342c30d5e8c", size = 470613 },
+    { url = "https://files.pythonhosted.org/packages/5d/4e/4e0bb5579b01fdbfd4388bd1eb9394a989e1336203a4b7f700d887b233c1/cffi-1.15.1-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:91fc98adde3d7881af9b59ed0294046f3806221863722ba7d8d120c575314325", size = 472199 },
+    { url = "https://files.pythonhosted.org/packages/37/5a/c37631a86be838bdd84cc0259130942bf7e6e32f70f4cab95f479847fb91/cffi-1.15.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:94411f22c3985acaec6f83c6df553f2dbe17b698cc7f8ae751ff2237d96b9e3c", size = 462588 },
+    { url = "https://files.pythonhosted.org/packages/71/d7/0fe0d91b0bbf610fb7254bb164fa8931596e660d62e90fb6289b7ee27b09/cffi-1.15.1-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:03425bdae262c76aad70202debd780501fabeaca237cdfddc008987c0e0f59ef", size = 450543 },
+    { url = "https://files.pythonhosted.org/packages/d3/56/3e94aa719ae96eeda8b68b3ec6e347e0a23168c6841dc276ccdcdadc9f32/cffi-1.15.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:cc4d65aeeaa04136a12677d3dd0b1c0c94dc43abac5860ab33cceb42b801c1e8", size = 474253 },
+    { url = "https://files.pythonhosted.org/packages/87/ee/ddc23981fc0f5e7b5356e98884226bcb899f95ebaefc3e8e8b8742dd7e22/cffi-1.15.1-cp311-cp311-win32.whl", hash = "sha256:a0f100c8912c114ff53e1202d0078b425bee3649ae34d7b070e9697f93c5d52d", size = 170313 },
+    { url = "https://files.pythonhosted.org/packages/43/a0/cc7370ef72b6ee586369bacd3961089ab3d94ae712febf07a244f1448ffd/cffi-1.15.1-cp311-cp311-win_amd64.whl", hash = "sha256:04ed324bda3cda42b9b695d51bb7d54b680b9719cfab04227cdd1e04e5de3104", size = 179001 },
+    { url = "https://files.pythonhosted.org/packages/b5/7d/df6c088ef30e78a78b0c9cca6b904d5abb698afb5bc8f5191d529d83d667/cffi-1.15.1-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:198caafb44239b60e252492445da556afafc7d1e3ab7a1fb3f0584ef6d742375", size = 178906 },
+    { url = "https://files.pythonhosted.org/packages/c2/0b/3b09a755ddb977c167e6d209a7536f6ade43bb0654bad42e08df1406b8e4/cffi-1.15.1-cp37-cp37m-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:5ef34d190326c3b1f822a5b7a45f6c4535e2f47ed06fec77d3d799c450b2651e", size = 405992 },
+    { url = "https://files.pythonhosted.org/packages/5b/1a/e1ee5bed11d8b6540c05a8e3c32448832d775364d4461dd6497374533401/cffi-1.15.1-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8102eaf27e1e448db915d08afa8b41d6c7ca7a04b7d73af6514df10a3e74bd82", size = 435560 },
+    { url = "https://files.pythonhosted.org/packages/d3/e1/e55ca2e0dd446caa2cc8f73c2b98879c04a1f4064ac529e1836683ca58b8/cffi-1.15.1-cp37-cp37m-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:5df2768244d19ab7f60546d0c7c63ce1581f7af8b5de3eb3004b9b6fc8a9f84b", size = 435478 },
+    { url = "https://files.pythonhosted.org/packages/2e/7a/68c35c151e5b7a12650ecc12fdfb85211aa1da43e9924598451c4a0a3839/cffi-1.15.1-cp37-cp37m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:a8c4917bd7ad33e8eb21e9a5bbba979b49d9a97acb3a803092cbc1133e20343c", size = 430395 },
+    { url = "https://files.pythonhosted.org/packages/93/d0/2e2b27ea2f69b0ec9e481647822f8f77f5fc23faca2dd00d1ff009940eb7/cffi-1.15.1-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0e2642fe3142e4cc4af0799748233ad6da94c62a8bec3a6648bf8ee68b1c7426", size = 427911 },
+    { url = "https://files.pythonhosted.org/packages/50/34/4cc590ad600869502c9838b4824982c122179089ed6791a8b1c95f0ff55e/cffi-1.15.1-cp37-cp37m-win32.whl", hash = "sha256:e229a521186c75c8ad9490854fd8bbdd9a0c9aa3a524326b55be83b54d4e0ad9", size = 169721 },
+    { url = "https://files.pythonhosted.org/packages/32/2a/63cb8c07d151de92ff9d897b2eb27ba6a0e78dda8e4c5f70d7b8c16cd6a2/cffi-1.15.1-cp37-cp37m-win_amd64.whl", hash = "sha256:a0b71b1b8fbf2b96e41c4d990244165e2c9be83d54962a9a1d118fd8657d2045", size = 179301 },
+    { url = "https://files.pythonhosted.org/packages/87/4b/64e8bd9d15d6b22b6cb11997094fbe61edf453ea0a97c8675cb7d1c3f06f/cffi-1.15.1-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:320dab6e7cb2eacdf0e658569d2575c4dad258c0fcc794f46215e1e39f90f2c3", size = 178940 },
+    { url = "https://files.pythonhosted.org/packages/22/c6/df826563f55f7e9dd9a1d3617866282afa969fe0d57decffa1911f416ed8/cffi-1.15.1-cp38-cp38-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1e74c6b51a9ed6589199c787bf5f9875612ca4a8a0785fb2d4a84429badaf22a", size = 421947 },
+    { url = "https://files.pythonhosted.org/packages/c1/25/16a082701378170559bb1d0e9ef2d293cece8dc62913d79351beb34c5ddf/cffi-1.15.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a5c84c68147988265e60416b57fc83425a78058853509c1b0629c180094904a5", size = 449906 },
+    { url = "https://files.pythonhosted.org/packages/df/02/aef53d4aa43154b829e9707c8c60bab413cd21819c4a36b0d7aaa83e2a61/cffi-1.15.1-cp38-cp38-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:3b926aa83d1edb5aa5b427b4053dc420ec295a08e40911296b9eb1b6170f6cca", size = 451028 },
+    { url = "https://files.pythonhosted.org/packages/79/4b/33494eb0adbcd884656c48f6db0c98ad8a5c678fb8fb5ed41ab546b04d8c/cffi-1.15.1-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:87c450779d0914f2861b8526e035c5e6da0a3199d8f1add1a665e1cbc6fc6d02", size = 446520 },
+    { url = "https://files.pythonhosted.org/packages/b7/8b/06f30caa03b5b3ac006de4f93478dbd0239e2a16566d81a106c322dc4f79/cffi-1.15.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4f2c9f67e9821cad2e5f480bc8d83b8742896f1242dba247911072d4fa94c192", size = 442655 },
+    { url = "https://files.pythonhosted.org/packages/47/97/137f0e3d2304df2060abb872a5830af809d7559a5a4b6a295afb02728e65/cffi-1.15.1-cp38-cp38-win32.whl", hash = "sha256:8b7ee99e510d7b66cdb6c593f21c043c248537a32e0bedf02e01e9553a172314", size = 170236 },
+    { url = "https://files.pythonhosted.org/packages/c9/e3/0a52838832408cfbbf3a59cb19bcd17e64eb33795c9710ca7d29ae10b5b7/cffi-1.15.1-cp38-cp38-win_amd64.whl", hash = "sha256:00a9ed42e88df81ffae7a8ab6d9356b371399b91dbdf0c3cb1e84c03a13aceb5", size = 178835 },
+    { url = "https://files.pythonhosted.org/packages/18/8f/5ff70c7458d61fa8a9752e5ee9c9984c601b0060aae0c619316a1e1f1ee5/cffi-1.15.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:54a2db7b78338edd780e7ef7f9f6c442500fb0d41a5a4ea24fff1c929d5af585", size = 179249 },
+    { url = "https://files.pythonhosted.org/packages/3a/75/a162315adeaf47e94a3b7f886a8e31d77b9e525a387eef2d6f0efc96a7c8/cffi-1.15.1-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:fcd131dd944808b5bdb38e6f5b53013c5aa4f334c5cad0c72742f6eba4b73db0", size = 174297 },
+    { url = "https://files.pythonhosted.org/packages/85/1f/a3c533f8d377da5ca7edb4f580cc3edc1edbebc45fac8bb3ae60f1176629/cffi-1.15.1-cp39-cp39-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:7473e861101c9e72452f9bf8acb984947aa1661a7704553a9f6e4baa5ba64415", size = 420641 },
+    { url = "https://files.pythonhosted.org/packages/77/b7/d3618d612be01e184033eab90006f8ca5b5edafd17bf247439ea4e167d8a/cffi-1.15.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6c9a799e985904922a4d207a94eae35c78ebae90e128f0c4e521ce339396be9d", size = 448814 },
+    { url = "https://files.pythonhosted.org/packages/a9/ba/e082df21ebaa9cb29f2c4e1d7e49a29b90fcd667d43632c6674a16d65382/cffi-1.15.1-cp39-cp39-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:3bcde07039e586f91b45c88f8583ea7cf7a0770df3a1649627bf598332cb6984", size = 449647 },
+    { url = "https://files.pythonhosted.org/packages/af/cb/53b7bba75a18372d57113ba934b27d0734206c283c1dfcc172347fbd9f76/cffi-1.15.1-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:33ab79603146aace82c2427da5ca6e58f2b3f2fb5da893ceac0c42218a40be35", size = 445191 },
+    { url = "https://files.pythonhosted.org/packages/2d/86/3ca57cddfa0419f6a95d1c8478f8f622ba597e3581fd501bbb915b20eb75/cffi-1.15.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5d598b938678ebf3c67377cdd45e09d431369c3b1a5b331058c338e201f12b27", size = 441236 },
+    { url = "https://files.pythonhosted.org/packages/ad/26/7b3a73ab7d82a64664c7c4ea470e4ec4a3c73bb4f02575c543a41e272de5/cffi-1.15.1-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:db0fbb9c62743ce59a9ff687eb5f4afbe77e5e8403d6697f7446e5f609976f76", size = 433865 },
+    { url = "https://files.pythonhosted.org/packages/da/ff/ab939e2c7b3f40d851c0f7192c876f1910f3442080c9c846532993ec3cef/cffi-1.15.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:98d85c6a2bef81588d9227dde12db8a7f47f639f4a17c9ae08e773aa9c697bf3", size = 463090 },
+    { url = "https://files.pythonhosted.org/packages/c6/3d/dd085bb831b22ce4d0b7ba8550e6d78960f02f770bbd1314fea3580727f8/cffi-1.15.1-cp39-cp39-win32.whl", hash = "sha256:40f4774f5a9d4f5e344f31a32b5096977b5d48560c5592e2f3d2c4374bd543ee", size = 170382 },
+    { url = "https://files.pythonhosted.org/packages/a8/16/06b84a7063a4c0a2b081030fdd976022086da9c14e80a9ed4ba0183a98a9/cffi-1.15.1-cp39-cp39-win_amd64.whl", hash = "sha256:70df4e3b545a17496c9b3f41f5115e69a4f2e77e94e1d2a8e1070bc0c38c8a3c", size = 179079 },
+]
+
+[[package]]
+name = "cfgv"
+version = "3.3.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/c4/bf/d0d622b660d414a47dc7f0d303791a627663f554345b21250e39e7acb48b/cfgv-3.3.1.tar.gz", hash = "sha256:f5a830efb9ce7a445376bb66ec94c638a9787422f96264c98edc6bdeed8ab736", size = 7864 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/6d/82/0a0ebd35bae9981dea55c06f8e6aaf44a49171ad798795c72c6f64cba4c2/cfgv-3.3.1-py2.py3-none-any.whl", hash = "sha256:c6a0883f3917a037485059700b9e75da2464e6c27051014ad85ba6aaa5884426", size = 7312 },
+]
+
+[[package]]
+name = "charset-normalizer"
+version = "3.4.3"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/83/2d/5fd176ceb9b2fc619e63405525573493ca23441330fcdaee6bef9460e924/charset_normalizer-3.4.3.tar.gz", hash = "sha256:6fce4b8500244f6fcb71465d4a4930d132ba9ab8e71a7859e6a5d59851068d14", size = 122371 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/d6/98/f3b8013223728a99b908c9344da3aa04ee6e3fa235f19409033eda92fb78/charset_normalizer-3.4.3-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:fb7f67a1bfa6e40b438170ebdc8158b78dc465a5a67b6dde178a46987b244a72", size = 207695 },
+    { url = "https://files.pythonhosted.org/packages/21/40/5188be1e3118c82dcb7c2a5ba101b783822cfb413a0268ed3be0468532de/charset_normalizer-3.4.3-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cc9370a2da1ac13f0153780040f465839e6cccb4a1e44810124b4e22483c93fe", size = 147153 },
+    { url = "https://files.pythonhosted.org/packages/37/60/5d0d74bc1e1380f0b72c327948d9c2aca14b46a9efd87604e724260f384c/charset_normalizer-3.4.3-cp310-cp310-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:07a0eae9e2787b586e129fdcbe1af6997f8d0e5abaa0bc98c0e20e124d67e601", size = 160428 },
+    { url = "https://files.pythonhosted.org/packages/85/9a/d891f63722d9158688de58d050c59dc3da560ea7f04f4c53e769de5140f5/charset_normalizer-3.4.3-cp310-cp310-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:74d77e25adda8581ffc1c720f1c81ca082921329452eba58b16233ab1842141c", size = 157627 },
+    { url = "https://files.pythonhosted.org/packages/65/1a/7425c952944a6521a9cfa7e675343f83fd82085b8af2b1373a2409c683dc/charset_normalizer-3.4.3-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d0e909868420b7049dafd3a31d45125b31143eec59235311fc4c57ea26a4acd2", size = 152388 },
+    { url = "https://files.pythonhosted.org/packages/f0/c9/a2c9c2a355a8594ce2446085e2ec97fd44d323c684ff32042e2a6b718e1d/charset_normalizer-3.4.3-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:c6f162aabe9a91a309510d74eeb6507fab5fff92337a15acbe77753d88d9dcf0", size = 150077 },
+    { url = "https://files.pythonhosted.org/packages/3b/38/20a1f44e4851aa1c9105d6e7110c9d020e093dfa5836d712a5f074a12bf7/charset_normalizer-3.4.3-cp310-cp310-musllinux_1_2_ppc64le.whl", hash = "sha256:4ca4c094de7771a98d7fbd67d9e5dbf1eb73efa4f744a730437d8a3a5cf994f0", size = 161631 },
+    { url = "https://files.pythonhosted.org/packages/a4/fa/384d2c0f57edad03d7bec3ebefb462090d8905b4ff5a2d2525f3bb711fac/charset_normalizer-3.4.3-cp310-cp310-musllinux_1_2_s390x.whl", hash = "sha256:02425242e96bcf29a49711b0ca9f37e451da7c70562bc10e8ed992a5a7a25cc0", size = 159210 },
+    { url = "https://files.pythonhosted.org/packages/33/9e/eca49d35867ca2db336b6ca27617deed4653b97ebf45dfc21311ce473c37/charset_normalizer-3.4.3-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:78deba4d8f9590fe4dae384aeff04082510a709957e968753ff3c48399f6f92a", size = 153739 },
+    { url = "https://files.pythonhosted.org/packages/2a/91/26c3036e62dfe8de8061182d33be5025e2424002125c9500faff74a6735e/charset_normalizer-3.4.3-cp310-cp310-win32.whl", hash = "sha256:d79c198e27580c8e958906f803e63cddb77653731be08851c7df0b1a14a8fc0f", size = 99825 },
+    { url = "https://files.pythonhosted.org/packages/e2/c6/f05db471f81af1fa01839d44ae2a8bfeec8d2a8b4590f16c4e7393afd323/charset_normalizer-3.4.3-cp310-cp310-win_amd64.whl", hash = "sha256:c6e490913a46fa054e03699c70019ab869e990270597018cef1d8562132c2669", size = 107452 },
+    { url = "https://files.pythonhosted.org/packages/7f/b5/991245018615474a60965a7c9cd2b4efbaabd16d582a5547c47ee1c7730b/charset_normalizer-3.4.3-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:b256ee2e749283ef3ddcff51a675ff43798d92d746d1a6e4631bf8c707d22d0b", size = 204483 },
+    { url = "https://files.pythonhosted.org/packages/c7/2a/ae245c41c06299ec18262825c1569c5d3298fc920e4ddf56ab011b417efd/charset_normalizer-3.4.3-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:13faeacfe61784e2559e690fc53fa4c5ae97c6fcedb8eb6fb8d0a15b475d2c64", size = 145520 },
+    { url = "https://files.pythonhosted.org/packages/3a/a4/b3b6c76e7a635748c4421d2b92c7b8f90a432f98bda5082049af37ffc8e3/charset_normalizer-3.4.3-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:00237675befef519d9af72169d8604a067d92755e84fe76492fef5441db05b91", size = 158876 },
+    { url = "https://files.pythonhosted.org/packages/e2/e6/63bb0e10f90a8243c5def74b5b105b3bbbfb3e7bb753915fe333fb0c11ea/charset_normalizer-3.4.3-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:585f3b2a80fbd26b048a0be90c5aae8f06605d3c92615911c3a2b03a8a3b796f", size = 156083 },
+    { url = "https://files.pythonhosted.org/packages/87/df/b7737ff046c974b183ea9aa111b74185ac8c3a326c6262d413bd5a1b8c69/charset_normalizer-3.4.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0e78314bdc32fa80696f72fa16dc61168fda4d6a0c014e0380f9d02f0e5d8a07", size = 150295 },
+    { url = "https://files.pythonhosted.org/packages/61/f1/190d9977e0084d3f1dc169acd060d479bbbc71b90bf3e7bf7b9927dec3eb/charset_normalizer-3.4.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:96b2b3d1a83ad55310de8c7b4a2d04d9277d5591f40761274856635acc5fcb30", size = 148379 },
+    { url = "https://files.pythonhosted.org/packages/4c/92/27dbe365d34c68cfe0ca76f1edd70e8705d82b378cb54ebbaeabc2e3029d/charset_normalizer-3.4.3-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:939578d9d8fd4299220161fdd76e86c6a251987476f5243e8864a7844476ba14", size = 160018 },
+    { url = "https://files.pythonhosted.org/packages/99/04/baae2a1ea1893a01635d475b9261c889a18fd48393634b6270827869fa34/charset_normalizer-3.4.3-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:fd10de089bcdcd1be95a2f73dbe6254798ec1bda9f450d5828c96f93e2536b9c", size = 157430 },
+    { url = "https://files.pythonhosted.org/packages/2f/36/77da9c6a328c54d17b960c89eccacfab8271fdaaa228305330915b88afa9/charset_normalizer-3.4.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:1e8ac75d72fa3775e0b7cb7e4629cec13b7514d928d15ef8ea06bca03ef01cae", size = 151600 },
+    { url = "https://files.pythonhosted.org/packages/64/d4/9eb4ff2c167edbbf08cdd28e19078bf195762e9bd63371689cab5ecd3d0d/charset_normalizer-3.4.3-cp311-cp311-win32.whl", hash = "sha256:6cf8fd4c04756b6b60146d98cd8a77d0cdae0e1ca20329da2ac85eed779b6849", size = 99616 },
+    { url = "https://files.pythonhosted.org/packages/f4/9c/996a4a028222e7761a96634d1820de8a744ff4327a00ada9c8942033089b/charset_normalizer-3.4.3-cp311-cp311-win_amd64.whl", hash = "sha256:31a9a6f775f9bcd865d88ee350f0ffb0e25936a7f930ca98995c05abf1faf21c", size = 107108 },
+    { url = "https://files.pythonhosted.org/packages/e9/5e/14c94999e418d9b87682734589404a25854d5f5d0408df68bc15b6ff54bb/charset_normalizer-3.4.3-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:e28e334d3ff134e88989d90ba04b47d84382a828c061d0d1027b1b12a62b39b1", size = 205655 },
+    { url = "https://files.pythonhosted.org/packages/7d/a8/c6ec5d389672521f644505a257f50544c074cf5fc292d5390331cd6fc9c3/charset_normalizer-3.4.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0cacf8f7297b0c4fcb74227692ca46b4a5852f8f4f24b3c766dd94a1075c4884", size = 146223 },
+    { url = "https://files.pythonhosted.org/packages/fc/eb/a2ffb08547f4e1e5415fb69eb7db25932c52a52bed371429648db4d84fb1/charset_normalizer-3.4.3-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c6fd51128a41297f5409deab284fecbe5305ebd7e5a1f959bee1c054622b7018", size = 159366 },
+    { url = "https://files.pythonhosted.org/packages/82/10/0fd19f20c624b278dddaf83b8464dcddc2456cb4b02bb902a6da126b87a1/charset_normalizer-3.4.3-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:3cfb2aad70f2c6debfbcb717f23b7eb55febc0bb23dcffc0f076009da10c6392", size = 157104 },
+    { url = "https://files.pythonhosted.org/packages/16/ab/0233c3231af734f5dfcf0844aa9582d5a1466c985bbed6cedab85af9bfe3/charset_normalizer-3.4.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1606f4a55c0fd363d754049cdf400175ee96c992b1f8018b993941f221221c5f", size = 151830 },
+    { url = "https://files.pythonhosted.org/packages/ae/02/e29e22b4e02839a0e4a06557b1999d0a47db3567e82989b5bb21f3fbbd9f/charset_normalizer-3.4.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:027b776c26d38b7f15b26a5da1044f376455fb3766df8fc38563b4efbc515154", size = 148854 },
+    { url = "https://files.pythonhosted.org/packages/05/6b/e2539a0a4be302b481e8cafb5af8792da8093b486885a1ae4d15d452bcec/charset_normalizer-3.4.3-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:42e5088973e56e31e4fa58eb6bd709e42fc03799c11c42929592889a2e54c491", size = 160670 },
+    { url = "https://files.pythonhosted.org/packages/31/e7/883ee5676a2ef217a40ce0bffcc3d0dfbf9e64cbcfbdf822c52981c3304b/charset_normalizer-3.4.3-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:cc34f233c9e71701040d772aa7490318673aa7164a0efe3172b2981218c26d93", size = 158501 },
+    { url = "https://files.pythonhosted.org/packages/c1/35/6525b21aa0db614cf8b5792d232021dca3df7f90a1944db934efa5d20bb1/charset_normalizer-3.4.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:320e8e66157cc4e247d9ddca8e21f427efc7a04bbd0ac8a9faf56583fa543f9f", size = 153173 },
+    { url = "https://files.pythonhosted.org/packages/50/ee/f4704bad8201de513fdc8aac1cabc87e38c5818c93857140e06e772b5892/charset_normalizer-3.4.3-cp312-cp312-win32.whl", hash = "sha256:fb6fecfd65564f208cbf0fba07f107fb661bcd1a7c389edbced3f7a493f70e37", size = 99822 },
+    { url = "https://files.pythonhosted.org/packages/39/f5/3b3836ca6064d0992c58c7561c6b6eee1b3892e9665d650c803bd5614522/charset_normalizer-3.4.3-cp312-cp312-win_amd64.whl", hash = "sha256:86df271bf921c2ee3818f0522e9a5b8092ca2ad8b065ece5d7d9d0e9f4849bcc", size = 107543 },
+    { url = "https://files.pythonhosted.org/packages/65/ca/2135ac97709b400c7654b4b764daf5c5567c2da45a30cdd20f9eefe2d658/charset_normalizer-3.4.3-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:14c2a87c65b351109f6abfc424cab3927b3bdece6f706e4d12faaf3d52ee5efe", size = 205326 },
+    { url = "https://files.pythonhosted.org/packages/71/11/98a04c3c97dd34e49c7d247083af03645ca3730809a5509443f3c37f7c99/charset_normalizer-3.4.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:41d1fc408ff5fdfb910200ec0e74abc40387bccb3252f3f27c0676731df2b2c8", size = 146008 },
+    { url = "https://files.pythonhosted.org/packages/60/f5/4659a4cb3c4ec146bec80c32d8bb16033752574c20b1252ee842a95d1a1e/charset_normalizer-3.4.3-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:1bb60174149316da1c35fa5233681f7c0f9f514509b8e399ab70fea5f17e45c9", size = 159196 },
+    { url = "https://files.pythonhosted.org/packages/86/9e/f552f7a00611f168b9a5865a1414179b2c6de8235a4fa40189f6f79a1753/charset_normalizer-3.4.3-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:30d006f98569de3459c2fc1f2acde170b7b2bd265dc1943e87e1a4efe1b67c31", size = 156819 },
+    { url = "https://files.pythonhosted.org/packages/7e/95/42aa2156235cbc8fa61208aded06ef46111c4d3f0de233107b3f38631803/charset_normalizer-3.4.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:416175faf02e4b0810f1f38bcb54682878a4af94059a1cd63b8747244420801f", size = 151350 },
+    { url = "https://files.pythonhosted.org/packages/c2/a9/3865b02c56f300a6f94fc631ef54f0a8a29da74fb45a773dfd3dcd380af7/charset_normalizer-3.4.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:6aab0f181c486f973bc7262a97f5aca3ee7e1437011ef0c2ec04b5a11d16c927", size = 148644 },
+    { url = "https://files.pythonhosted.org/packages/77/d9/cbcf1a2a5c7d7856f11e7ac2d782aec12bdfea60d104e60e0aa1c97849dc/charset_normalizer-3.4.3-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:fdabf8315679312cfa71302f9bd509ded4f2f263fb5b765cf1433b39106c3cc9", size = 160468 },
+    { url = "https://files.pythonhosted.org/packages/f6/42/6f45efee8697b89fda4d50580f292b8f7f9306cb2971d4b53f8914e4d890/charset_normalizer-3.4.3-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:bd28b817ea8c70215401f657edef3a8aa83c29d447fb0b622c35403780ba11d5", size = 158187 },
+    { url = "https://files.pythonhosted.org/packages/70/99/f1c3bdcfaa9c45b3ce96f70b14f070411366fa19549c1d4832c935d8e2c3/charset_normalizer-3.4.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:18343b2d246dc6761a249ba1fb13f9ee9a2bcd95decc767319506056ea4ad4dc", size = 152699 },
+    { url = "https://files.pythonhosted.org/packages/a3/ad/b0081f2f99a4b194bcbb1934ef3b12aa4d9702ced80a37026b7607c72e58/charset_normalizer-3.4.3-cp313-cp313-win32.whl", hash = "sha256:6fb70de56f1859a3f71261cbe41005f56a7842cc348d3aeb26237560bfa5e0ce", size = 99580 },
+    { url = "https://files.pythonhosted.org/packages/9a/8f/ae790790c7b64f925e5c953b924aaa42a243fb778fed9e41f147b2a5715a/charset_normalizer-3.4.3-cp313-cp313-win_amd64.whl", hash = "sha256:cf1ebb7d78e1ad8ec2a8c4732c7be2e736f6e5123a4146c5b89c9d1f585f8cef", size = 107366 },
+    { url = "https://files.pythonhosted.org/packages/8e/91/b5a06ad970ddc7a0e513112d40113e834638f4ca1120eb727a249fb2715e/charset_normalizer-3.4.3-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:3cd35b7e8aedeb9e34c41385fda4f73ba609e561faedfae0a9e75e44ac558a15", size = 204342 },
+    { url = "https://files.pythonhosted.org/packages/ce/ec/1edc30a377f0a02689342f214455c3f6c2fbedd896a1d2f856c002fc3062/charset_normalizer-3.4.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b89bc04de1d83006373429975f8ef9e7932534b8cc9ca582e4db7d20d91816db", size = 145995 },
+    { url = "https://files.pythonhosted.org/packages/17/e5/5e67ab85e6d22b04641acb5399c8684f4d37caf7558a53859f0283a650e9/charset_normalizer-3.4.3-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:2001a39612b241dae17b4687898843f254f8748b796a2e16f1051a17078d991d", size = 158640 },
+    { url = "https://files.pythonhosted.org/packages/f1/e5/38421987f6c697ee3722981289d554957c4be652f963d71c5e46a262e135/charset_normalizer-3.4.3-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:8dcfc373f888e4fb39a7bc57e93e3b845e7f462dacc008d9749568b1c4ece096", size = 156636 },
+    { url = "https://files.pythonhosted.org/packages/a0/e4/5a075de8daa3ec0745a9a3b54467e0c2967daaaf2cec04c845f73493e9a1/charset_normalizer-3.4.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:18b97b8404387b96cdbd30ad660f6407799126d26a39ca65729162fd810a99aa", size = 150939 },
+    { url = "https://files.pythonhosted.org/packages/02/f7/3611b32318b30974131db62b4043f335861d4d9b49adc6d57c1149cc49d4/charset_normalizer-3.4.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:ccf600859c183d70eb47e05a44cd80a4ce77394d1ac0f79dbd2dd90a69a3a049", size = 148580 },
+    { url = "https://files.pythonhosted.org/packages/7e/61/19b36f4bd67f2793ab6a99b979b4e4f3d8fc754cbdffb805335df4337126/charset_normalizer-3.4.3-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:53cd68b185d98dde4ad8990e56a58dea83a4162161b1ea9272e5c9182ce415e0", size = 159870 },
+    { url = "https://files.pythonhosted.org/packages/06/57/84722eefdd338c04cf3030ada66889298eaedf3e7a30a624201e0cbe424a/charset_normalizer-3.4.3-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:30a96e1e1f865f78b030d65241c1ee850cdf422d869e9028e2fc1d5e4db73b92", size = 157797 },
+    { url = "https://files.pythonhosted.org/packages/72/2a/aff5dd112b2f14bcc3462c312dce5445806bfc8ab3a7328555da95330e4b/charset_normalizer-3.4.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:d716a916938e03231e86e43782ca7878fb602a125a91e7acb8b5112e2e96ac16", size = 152224 },
+    { url = "https://files.pythonhosted.org/packages/b7/8c/9839225320046ed279c6e839d51f028342eb77c91c89b8ef2549f951f3ec/charset_normalizer-3.4.3-cp314-cp314-win32.whl", hash = "sha256:c6dbd0ccdda3a2ba7c2ecd9d77b37f3b5831687d8dc1b6ca5f56a4880cc7b7ce", size = 100086 },
+    { url = "https://files.pythonhosted.org/packages/ee/7a/36fbcf646e41f710ce0a563c1c9a343c6edf9be80786edeb15b6f62e17db/charset_normalizer-3.4.3-cp314-cp314-win_amd64.whl", hash = "sha256:73dc19b562516fc9bcf6e5d6e596df0b4eb98d87e4f79f3ae71840e6ed21361c", size = 107400 },
+    { url = "https://files.pythonhosted.org/packages/22/82/63a45bfc36f73efe46731a3a71cb84e2112f7e0b049507025ce477f0f052/charset_normalizer-3.4.3-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:0f2be7e0cf7754b9a30eb01f4295cc3d4358a479843b31f328afd210e2c7598c", size = 198805 },
+    { url = "https://files.pythonhosted.org/packages/0c/52/8b0c6c3e53f7e546a5e49b9edb876f379725914e1130297f3b423c7b71c5/charset_normalizer-3.4.3-cp38-cp38-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c60e092517a73c632ec38e290eba714e9627abe9d301c8c8a12ec32c314a2a4b", size = 142862 },
+    { url = "https://files.pythonhosted.org/packages/59/c0/a74f3bd167d311365e7973990243f32c35e7a94e45103125275b9e6c479f/charset_normalizer-3.4.3-cp38-cp38-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:252098c8c7a873e17dd696ed98bbe91dbacd571da4b87df3736768efa7a792e4", size = 155104 },
+    { url = "https://files.pythonhosted.org/packages/1a/79/ae516e678d6e32df2e7e740a7be51dc80b700e2697cb70054a0f1ac2c955/charset_normalizer-3.4.3-cp38-cp38-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:3653fad4fe3ed447a596ae8638b437f827234f01a8cd801842e43f3d0a6b281b", size = 152598 },
+    { url = "https://files.pythonhosted.org/packages/00/bd/ef9c88464b126fa176f4ef4a317ad9b6f4d30b2cffbc43386062367c3e2c/charset_normalizer-3.4.3-cp38-cp38-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8999f965f922ae054125286faf9f11bc6932184b93011d138925a1773830bbe9", size = 147391 },
+    { url = "https://files.pythonhosted.org/packages/7a/03/cbb6fac9d3e57f7e07ce062712ee80d80a5ab46614684078461917426279/charset_normalizer-3.4.3-cp38-cp38-musllinux_1_2_aarch64.whl", hash = "sha256:d95bfb53c211b57198bb91c46dd5a2d8018b3af446583aab40074bf7988401cb", size = 145037 },
+    { url = "https://files.pythonhosted.org/packages/64/d1/f9d141c893ef5d4243bc75c130e95af8fd4bc355beff06e9b1e941daad6e/charset_normalizer-3.4.3-cp38-cp38-musllinux_1_2_ppc64le.whl", hash = "sha256:5b413b0b1bfd94dbf4023ad6945889f374cd24e3f62de58d6bb102c4d9ae534a", size = 156425 },
+    { url = "https://files.pythonhosted.org/packages/c5/35/9c99739250742375167bc1b1319cd1cec2bf67438a70d84b2e1ec4c9daa3/charset_normalizer-3.4.3-cp38-cp38-musllinux_1_2_s390x.whl", hash = "sha256:b5e3b2d152e74e100a9e9573837aba24aab611d39428ded46f4e4022ea7d1942", size = 153734 },
+    { url = "https://files.pythonhosted.org/packages/50/10/c117806094d2c956ba88958dab680574019abc0c02bcf57b32287afca544/charset_normalizer-3.4.3-cp38-cp38-musllinux_1_2_x86_64.whl", hash = "sha256:a2d08ac246bb48479170408d6c19f6385fa743e7157d716e144cad849b2dd94b", size = 148551 },
+    { url = "https://files.pythonhosted.org/packages/61/c5/dc3ba772489c453621ffc27e8978a98fe7e41a93e787e5e5bde797f1dddb/charset_normalizer-3.4.3-cp38-cp38-win32.whl", hash = "sha256:ec557499516fc90fd374bf2e32349a2887a876fbf162c160e3c01b6849eaf557", size = 98459 },
+    { url = "https://files.pythonhosted.org/packages/05/35/bb59b1cd012d7196fc81c2f5879113971efc226a63812c9cf7f89fe97c40/charset_normalizer-3.4.3-cp38-cp38-win_amd64.whl", hash = "sha256:5d8d01eac18c423815ed4f4a2ec3b439d654e55ee4ad610e153cf02faf67ea40", size = 105887 },
+    { url = "https://files.pythonhosted.org/packages/c2/ca/9a0983dd5c8e9733565cf3db4df2b0a2e9a82659fd8aa2a868ac6e4a991f/charset_normalizer-3.4.3-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:70bfc5f2c318afece2f5838ea5e4c3febada0be750fcf4775641052bbba14d05", size = 207520 },
+    { url = "https://files.pythonhosted.org/packages/39/c6/99271dc37243a4f925b09090493fb96c9333d7992c6187f5cfe5312008d2/charset_normalizer-3.4.3-cp39-cp39-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:23b6b24d74478dc833444cbd927c338349d6ae852ba53a0d02a2de1fce45b96e", size = 147307 },
+    { url = "https://files.pythonhosted.org/packages/e4/69/132eab043356bba06eb333cc2cc60c6340857d0a2e4ca6dc2b51312886b3/charset_normalizer-3.4.3-cp39-cp39-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:34a7f768e3f985abdb42841e20e17b330ad3aaf4bb7e7aeeb73db2e70f077b99", size = 160448 },
+    { url = "https://files.pythonhosted.org/packages/04/9a/914d294daa4809c57667b77470533e65def9c0be1ef8b4c1183a99170e9d/charset_normalizer-3.4.3-cp39-cp39-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:fb731e5deb0c7ef82d698b0f4c5bb724633ee2a489401594c5c88b02e6cb15f7", size = 157758 },
+    { url = "https://files.pythonhosted.org/packages/b0/a8/6f5bcf1bcf63cb45625f7c5cadca026121ff8a6c8a3256d8d8cd59302663/charset_normalizer-3.4.3-cp39-cp39-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:257f26fed7d7ff59921b78244f3cd93ed2af1800ff048c33f624c87475819dd7", size = 152487 },
+    { url = "https://files.pythonhosted.org/packages/c4/72/d3d0e9592f4e504f9dea08b8db270821c909558c353dc3b457ed2509f2fb/charset_normalizer-3.4.3-cp39-cp39-musllinux_1_2_aarch64.whl", hash = "sha256:1ef99f0456d3d46a50945c98de1774da86f8e992ab5c77865ea8b8195341fc19", size = 150054 },
+    { url = "https://files.pythonhosted.org/packages/20/30/5f64fe3981677fe63fa987b80e6c01042eb5ff653ff7cec1b7bd9268e54e/charset_normalizer-3.4.3-cp39-cp39-musllinux_1_2_ppc64le.whl", hash = "sha256:2c322db9c8c89009a990ef07c3bcc9f011a3269bc06782f916cd3d9eed7c9312", size = 161703 },
+    { url = "https://files.pythonhosted.org/packages/e1/ef/dd08b2cac9284fd59e70f7d97382c33a3d0a926e45b15fc21b3308324ffd/charset_normalizer-3.4.3-cp39-cp39-musllinux_1_2_s390x.whl", hash = "sha256:511729f456829ef86ac41ca78c63a5cb55240ed23b4b737faca0eb1abb1c41bc", size = 159096 },
+    { url = "https://files.pythonhosted.org/packages/45/8c/dcef87cfc2b3f002a6478f38906f9040302c68aebe21468090e39cde1445/charset_normalizer-3.4.3-cp39-cp39-musllinux_1_2_x86_64.whl", hash = "sha256:88ab34806dea0671532d3f82d82b85e8fc23d7b2dd12fa837978dad9bb392a34", size = 153852 },
+    { url = "https://files.pythonhosted.org/packages/63/86/9cbd533bd37883d467fcd1bd491b3547a3532d0fbb46de2b99feeebf185e/charset_normalizer-3.4.3-cp39-cp39-win32.whl", hash = "sha256:16a8770207946ac75703458e2c743631c79c59c5890c80011d536248f8eaa432", size = 99840 },
+    { url = "https://files.pythonhosted.org/packages/ce/d6/7e805c8e5c46ff9729c49950acc4ee0aeb55efb8b3a56687658ad10c3216/charset_normalizer-3.4.3-cp39-cp39-win_amd64.whl", hash = "sha256:d22dbedd33326a4a5190dd4fe9e9e693ef12160c77382d9e87919bce54f3d4ca", size = 107438 },
+    { url = "https://files.pythonhosted.org/packages/8a/1f/f041989e93b001bc4e44bb1669ccdcf54d3f00e628229a85b08d330615c5/charset_normalizer-3.4.3-py3-none-any.whl", hash = "sha256:ce571ab16d890d23b5c278547ba694193a45011ff86a9162a71307ed9f86759a", size = 53175 },
+]
+
+[[package]]
+name = "click"
+version = "8.1.8"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "colorama", marker = "platform_system == 'Windows'" },
+    { name = "importlib-metadata", marker = "python_full_version < '3.8'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/b9/2e/0090cbf739cee7d23781ad4b89a9894a41538e4fcf4c31dcdd705b78eb8b/click-8.1.8.tar.gz", hash = "sha256:ed53c9d8990d83c2a27deae68e4ee337473f6330c040a31d4225c9574d16096a", size = 226593 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/7e/d4/7ebdbd03970677812aac39c869717059dbb71a4cfc033ca6e5221787892c/click-8.1.8-py3-none-any.whl", hash = "sha256:63c132bbbed01578a06712a2d1f497bb62d9c1c0d329b7903a866228027263b2", size = 98188 },
+]
+
+[[package]]
+name = "colorama"
+version = "0.4.6"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/d8/53/6f443c9a4a8358a93a6792e2acffb9d9d5cb0a5cfd8802644b7b1c9a02e4/colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44", size = 27697 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335 },
+]
+
+[[package]]
+name = "distlib"
+version = "0.4.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/96/8e/709914eb2b5749865801041647dc7f4e6d00b549cfe88b65ca192995f07c/distlib-0.4.0.tar.gz", hash = "sha256:feec40075be03a04501a973d81f633735b4b69f98b05450592310c0f401a4e0d", size = 614605 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/33/6b/e0547afaf41bf2c42e52430072fa5658766e3d65bd4b03a563d1b6336f57/distlib-0.4.0-py2.py3-none-any.whl", hash = "sha256:9659f7d87e46584a30b5780e43ac7a2143098441670ff0a49d5f9034c54a6c16", size = 469047 },
+]
+
+[[package]]
+name = "exceptiongroup"
+version = "1.3.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "typing-extensions", marker = "python_full_version < '3.13'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/0b/9f/a65090624ecf468cdca03533906e7c69ed7588582240cfe7cc9e770b50eb/exceptiongroup-1.3.0.tar.gz", hash = "sha256:b241f5885f560bc56a59ee63ca4c6a8bfa46ae4ad651af316d4e81817bb9fd88", size = 29749 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/36/f4/c6e662dade71f56cd2f3735141b265c3c79293c109549c1e6933b0651ffc/exceptiongroup-1.3.0-py3-none-any.whl", hash = "sha256:4d111e6e0c13d0644cad6ddaa7ed0261a0b36971f6d23e7ec9b4b9097da78a10", size = 16674 },
+]
+
+[[package]]
+name = "filelock"
+version = "3.12.2"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/00/0b/c506e9e44e4c4b6c89fcecda23dc115bf8e7ff7eb127e0cb9c114cbc9a15/filelock-3.12.2.tar.gz", hash = "sha256:002740518d8aa59a26b0c76e10fb8c6e15eae825d34b6fdf670333fd7b938d81", size = 12441 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/00/45/ec3407adf6f6b5bf867a4462b2b0af27597a26bd3cd6e2534cb6ab029938/filelock-3.12.2-py3-none-any.whl", hash = "sha256:cbb791cdea2a72f23da6ac5b5269ab0a0d161e9ef0100e653b69049a7706d1ec", size = 10923 },
+]
+
+[[package]]
+name = "flask"
+version = "2.2.5"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "click" },
+    { name = "importlib-metadata", marker = "python_full_version < '3.10'" },
+    { name = "itsdangerous" },
+    { name = "jinja2" },
+    { name = "werkzeug" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/5f/76/a4d2c4436dda4b0a12c71e075c508ea7988a1066b06a575f6afe4fecc023/Flask-2.2.5.tar.gz", hash = "sha256:edee9b0a7ff26621bd5a8c10ff484ae28737a2410d99b0bb9a6850c7fb977aa0", size = 697814 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/9f/1a/8b6d48162861009d1e017a9740431c78d860809773b66cac220a11aa3310/Flask-2.2.5-py3-none-any.whl", hash = "sha256:58107ed83443e86067e41eff4631b058178191a355886f8e479e347fa1285fdf", size = 101817 },
+]
+
+[[package]]
+name = "flask-cloudflared"
+version = "0.0.14"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "flask" },
+    { name = "requests" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/49/cf/d01e9f44c042bb77194c83d00bc98bcf1cd8ff24cdf686191d8d57b8d4dc/flask_cloudflared-0.0.14.tar.gz", hash = "sha256:be88151714f306d042923c6996332de39519c24a1914692547d92789d098912a", size = 6157 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/f9/e0/d90aa747ae35adfcd6fa116936b1eb42a6e5d8c69419694862476cefd34b/flask_cloudflared-0.0.14-py3-none-any.whl", hash = "sha256:abd11c16030d5c68f91cf64efaecdebf09ab0dce10db43333aaebc790569bdbf", size = 6420 },
+]
+
+[[package]]
+name = "identify"
+version = "2.5.24"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/c4/f8/498e13e408d25ee6ff04aa0acbf91ad8e9caae74be91720fc0e811e649b7/identify-2.5.24.tar.gz", hash = "sha256:0aac67d5b4812498056d28a9a512a483f5085cc28640b02b258a59dac34301d4", size = 98886 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/4f/fd/2c46fba2bc032ba4c970bb8de59d25187087d7138a0ebf7c1dcc91d94f01/identify-2.5.24-py2.py3-none-any.whl", hash = "sha256:986dbfb38b1140e763e413e6feb44cd731faf72d1909543178aa79b0e258265d", size = 98826 },
+]
+
+[[package]]
+name = "idna"
+version = "3.10"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/f1/70/7703c29685631f5a7590aa73f1f1d3fa9a380e654b86af429e0934a32f7d/idna-3.10.tar.gz", hash = "sha256:12f65c9b470abda6dc35cf8e63cc574b1c52b11df2c86030af0ac09b01b13ea9", size = 190490 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/76/c6/c88e154df9c4e1a2a66ccf0005a88dfb2650c1dffb6f5ce603dfbd452ce3/idna-3.10-py3-none-any.whl", hash = "sha256:946d195a0d259cbba61165e88e65941f16e9b36ea6ddb97f00452bae8b1287d3", size = 70442 },
+]
+
+[[package]]
+name = "importlib-metadata"
+version = "6.7.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "typing-extensions", marker = "python_full_version < '3.8'" },
+    { name = "zipp" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/a3/82/f6e29c8d5c098b6be61460371c2c5591f4a335923639edec43b3830650a4/importlib_metadata-6.7.0.tar.gz", hash = "sha256:1aaf550d4f73e5d6783e7acb77aec43d49da8017410afae93822cc9cca98c4d4", size = 53569 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/ff/94/64287b38c7de4c90683630338cf28f129decbba0a44f0c6db35a873c73c4/importlib_metadata-6.7.0-py3-none-any.whl", hash = "sha256:cb52082e659e97afc5dac71e79de97d8681de3aa07ff18578330904a9d18e5b5", size = 22934 },
+]
+
+[[package]]
+name = "importlib-resources"
+version = "5.12.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "zipp", marker = "python_full_version < '3.10'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/4e/a2/3cab1de83f95dd15297c15bdc04d50902391d707247cada1f021bbfe2149/importlib_resources-5.12.0.tar.gz", hash = "sha256:4be82589bf5c1d7999aedf2a45159d10cb3ca4f19b2271f8792bc8e6da7b22f6", size = 39894 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/38/71/c13ea695a4393639830bf96baea956538ba7a9d06fcce7cef10bfff20f72/importlib_resources-5.12.0-py3-none-any.whl", hash = "sha256:7b1deeebbf351c7578e09bf2f63fa2ce8b5ffec296e0d349139d43cca061a81a", size = 36211 },
+]
+
+[[package]]
+name = "iniconfig"
+version = "2.0.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/d7/4b/cbd8e699e64a6f16ca3a8220661b5f83792b3017d0f79807cb8708d33913/iniconfig-2.0.0.tar.gz", hash = "sha256:2d91e135bf72d31a410b17c16da610a82cb55f6b0477d1a902134b24a455b8b3", size = 4646 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/ef/a6/62565a6e1cf69e10f5727360368e451d4b7f58beeac6173dc9db836a5b46/iniconfig-2.0.0-py3-none-any.whl", hash = "sha256:b6a85871a79d2e3b22d2d1b94ac2824226a63c6b741c88f7ae975f18b6778374", size = 5892 },
+]
+
+[[package]]
+name = "itsdangerous"
+version = "2.1.2"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/7f/a1/d3fb83e7a61fa0c0d3d08ad0a94ddbeff3731c05212617dff3a94e097f08/itsdangerous-2.1.2.tar.gz", hash = "sha256:5dbbc68b317e5e42f327f9021763545dc3fc3bfe22e6deb96aaf1fc38874156a", size = 56143 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/68/5f/447e04e828f47465eeab35b5d408b7ebaaaee207f48b7136c5a7267a30ae/itsdangerous-2.1.2-py3-none-any.whl", hash = "sha256:2c2349112351b88699d8d4b6b075022c0808887cb7ad10069318a8b0bc88db44", size = 15749 },
+]
+
+[[package]]
+name = "jinja2"
+version = "3.1.6"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "markupsafe" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/df/bf/f7da0350254c0ed7c72f3e33cef02e048281fec7ecec5f032d4aac52226b/jinja2-3.1.6.tar.gz", hash = "sha256:0137fb05990d35f1275a587e9aee6d56da821fc83491a0fb838183be43f66d6d", size = 245115 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899 },
+]
+
+[[package]]
+name = "markupsafe"
+version = "2.1.5"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/87/5b/aae44c6655f3801e81aa3eef09dbbf012431987ba564d7231722f68df02d/MarkupSafe-2.1.5.tar.gz", hash = "sha256:d283d37a890ba4c1ae73ffadf8046435c76e7bc2247bbb63c00bd1a709c6544b", size = 19384 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/e4/54/ad5eb37bf9d51800010a74e4665425831a9db4e7c4e0fde4352e391e808e/MarkupSafe-2.1.5-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:a17a92de5231666cfbe003f0e4b9b3a7ae3afb1ec2845aadc2bacc93ff85febc", size = 18206 },
+    { url = "https://files.pythonhosted.org/packages/6a/4a/a4d49415e600bacae038c67f9fecc1d5433b9d3c71a4de6f33537b89654c/MarkupSafe-2.1.5-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:72b6be590cc35924b02c78ef34b467da4ba07e4e0f0454a2c5907f473fc50ce5", size = 14079 },
+    { url = "https://files.pythonhosted.org/packages/0a/7b/85681ae3c33c385b10ac0f8dd025c30af83c78cec1c37a6aa3b55e67f5ec/MarkupSafe-2.1.5-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e61659ba32cf2cf1481e575d0462554625196a1f2fc06a1c777d3f48e8865d46", size = 26620 },
+    { url = "https://files.pythonhosted.org/packages/7c/52/2b1b570f6b8b803cef5ac28fdf78c0da318916c7d2fe9402a84d591b394c/MarkupSafe-2.1.5-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2174c595a0d73a3080ca3257b40096db99799265e1c27cc5a610743acd86d62f", size = 25818 },
+    { url = "https://files.pythonhosted.org/packages/29/fe/a36ba8c7ca55621620b2d7c585313efd10729e63ef81e4e61f52330da781/MarkupSafe-2.1.5-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ae2ad8ae6ebee9d2d94b17fb62763125f3f374c25618198f40cbb8b525411900", size = 25493 },
+    { url = "https://files.pythonhosted.org/packages/60/ae/9c60231cdfda003434e8bd27282b1f4e197ad5a710c14bee8bea8a9ca4f0/MarkupSafe-2.1.5-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:075202fa5b72c86ad32dc7d0b56024ebdbcf2048c0ba09f1cde31bfdd57bcfff", size = 30630 },
+    { url = "https://files.pythonhosted.org/packages/65/dc/1510be4d179869f5dafe071aecb3f1f41b45d37c02329dfba01ff59e5ac5/MarkupSafe-2.1.5-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:598e3276b64aff0e7b3451b72e94fa3c238d452e7ddcd893c3ab324717456bad", size = 29745 },
+    { url = "https://files.pythonhosted.org/packages/30/39/8d845dd7d0b0613d86e0ef89549bfb5f61ed781f59af45fc96496e897f3a/MarkupSafe-2.1.5-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:fce659a462a1be54d2ffcacea5e3ba2d74daa74f30f5f143fe0c58636e355fdd", size = 30021 },
+    { url = "https://files.pythonhosted.org/packages/c7/5c/356a6f62e4f3c5fbf2602b4771376af22a3b16efa74eb8716fb4e328e01e/MarkupSafe-2.1.5-cp310-cp310-win32.whl", hash = "sha256:d9fad5155d72433c921b782e58892377c44bd6252b5af2f67f16b194987338a4", size = 16659 },
+    { url = "https://files.pythonhosted.org/packages/69/48/acbf292615c65f0604a0c6fc402ce6d8c991276e16c80c46a8f758fbd30c/MarkupSafe-2.1.5-cp310-cp310-win_amd64.whl", hash = "sha256:bf50cd79a75d181c9181df03572cdce0fbb75cc353bc350712073108cba98de5", size = 17213 },
+    { url = "https://files.pythonhosted.org/packages/11/e7/291e55127bb2ae67c64d66cef01432b5933859dfb7d6949daa721b89d0b3/MarkupSafe-2.1.5-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:629ddd2ca402ae6dbedfceeba9c46d5f7b2a61d9749597d4307f943ef198fc1f", size = 18219 },
+    { url = "https://files.pythonhosted.org/packages/6b/cb/aed7a284c00dfa7c0682d14df85ad4955a350a21d2e3b06d8240497359bf/MarkupSafe-2.1.5-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:5b7b716f97b52c5a14bffdf688f971b2d5ef4029127f1ad7a513973cfd818df2", size = 14098 },
+    { url = "https://files.pythonhosted.org/packages/1c/cf/35fe557e53709e93feb65575c93927942087e9b97213eabc3fe9d5b25a55/MarkupSafe-2.1.5-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6ec585f69cec0aa07d945b20805be741395e28ac1627333b1c5b0105962ffced", size = 29014 },
+    { url = "https://files.pythonhosted.org/packages/97/18/c30da5e7a0e7f4603abfc6780574131221d9148f323752c2755d48abad30/MarkupSafe-2.1.5-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b91c037585eba9095565a3556f611e3cbfaa42ca1e865f7b8015fe5c7336d5a5", size = 28220 },
+    { url = "https://files.pythonhosted.org/packages/0c/40/2e73e7d532d030b1e41180807a80d564eda53babaf04d65e15c1cf897e40/MarkupSafe-2.1.5-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:7502934a33b54030eaf1194c21c692a534196063db72176b0c4028e140f8f32c", size = 27756 },
+    { url = "https://files.pythonhosted.org/packages/18/46/5dca760547e8c59c5311b332f70605d24c99d1303dd9a6e1fc3ed0d73561/MarkupSafe-2.1.5-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:0e397ac966fdf721b2c528cf028494e86172b4feba51d65f81ffd65c63798f3f", size = 33988 },
+    { url = "https://files.pythonhosted.org/packages/6d/c5/27febe918ac36397919cd4a67d5579cbbfa8da027fa1238af6285bb368ea/MarkupSafe-2.1.5-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:c061bb86a71b42465156a3ee7bd58c8c2ceacdbeb95d05a99893e08b8467359a", size = 32718 },
+    { url = "https://files.pythonhosted.org/packages/f8/81/56e567126a2c2bc2684d6391332e357589a96a76cb9f8e5052d85cb0ead8/MarkupSafe-2.1.5-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:3a57fdd7ce31c7ff06cdfbf31dafa96cc533c21e443d57f5b1ecc6cdc668ec7f", size = 33317 },
+    { url = "https://files.pythonhosted.org/packages/00/0b/23f4b2470accb53285c613a3ab9ec19dc944eaf53592cb6d9e2af8aa24cc/MarkupSafe-2.1.5-cp311-cp311-win32.whl", hash = "sha256:397081c1a0bfb5124355710fe79478cdbeb39626492b15d399526ae53422b906", size = 16670 },
+    { url = "https://files.pythonhosted.org/packages/b7/a2/c78a06a9ec6d04b3445a949615c4c7ed86a0b2eb68e44e7541b9d57067cc/MarkupSafe-2.1.5-cp311-cp311-win_amd64.whl", hash = "sha256:2b7c57a4dfc4f16f7142221afe5ba4e093e09e728ca65c51f5620c9aaeb9a617", size = 17224 },
+    { url = "https://files.pythonhosted.org/packages/53/bd/583bf3e4c8d6a321938c13f49d44024dbe5ed63e0a7ba127e454a66da974/MarkupSafe-2.1.5-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:8dec4936e9c3100156f8a2dc89c4b88d5c435175ff03413b443469c7c8c5f4d1", size = 18215 },
+    { url = "https://files.pythonhosted.org/packages/48/d6/e7cd795fc710292c3af3a06d80868ce4b02bfbbf370b7cee11d282815a2a/MarkupSafe-2.1.5-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:3c6b973f22eb18a789b1460b4b91bf04ae3f0c4234a0a6aa6b0a92f6f7b951d4", size = 14069 },
+    { url = "https://files.pythonhosted.org/packages/51/b5/5d8ec796e2a08fc814a2c7d2584b55f889a55cf17dd1a90f2beb70744e5c/MarkupSafe-2.1.5-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ac07bad82163452a6884fe8fa0963fb98c2346ba78d779ec06bd7a6262132aee", size = 29452 },
+    { url = "https://files.pythonhosted.org/packages/0a/0d/2454f072fae3b5a137c119abf15465d1771319dfe9e4acbb31722a0fff91/MarkupSafe-2.1.5-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f5dfb42c4604dddc8e4305050aa6deb084540643ed5804d7455b5df8fe16f5e5", size = 28462 },
+    { url = "https://files.pythonhosted.org/packages/2d/75/fd6cb2e68780f72d47e6671840ca517bda5ef663d30ada7616b0462ad1e3/MarkupSafe-2.1.5-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ea3d8a3d18833cf4304cd2fc9cbb1efe188ca9b5efef2bdac7adc20594a0e46b", size = 27869 },
+    { url = "https://files.pythonhosted.org/packages/b0/81/147c477391c2750e8fc7705829f7351cf1cd3be64406edcf900dc633feb2/MarkupSafe-2.1.5-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:d050b3361367a06d752db6ead6e7edeb0009be66bc3bae0ee9d97fb326badc2a", size = 33906 },
+    { url = "https://files.pythonhosted.org/packages/8b/ff/9a52b71839d7a256b563e85d11050e307121000dcebc97df120176b3ad93/MarkupSafe-2.1.5-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:bec0a414d016ac1a18862a519e54b2fd0fc8bbfd6890376898a6c0891dd82e9f", size = 32296 },
+    { url = "https://files.pythonhosted.org/packages/88/07/2dc76aa51b481eb96a4c3198894f38b480490e834479611a4053fbf08623/MarkupSafe-2.1.5-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:58c98fee265677f63a4385256a6d7683ab1832f3ddd1e66fe948d5880c21a169", size = 33038 },
+    { url = "https://files.pythonhosted.org/packages/96/0c/620c1fb3661858c0e37eb3cbffd8c6f732a67cd97296f725789679801b31/MarkupSafe-2.1.5-cp312-cp312-win32.whl", hash = "sha256:8590b4ae07a35970728874632fed7bd57b26b0102df2d2b233b6d9d82f6c62ad", size = 16572 },
+    { url = "https://files.pythonhosted.org/packages/3f/14/c3554d512d5f9100a95e737502f4a2323a1959f6d0d01e0d0997b35f7b10/MarkupSafe-2.1.5-cp312-cp312-win_amd64.whl", hash = "sha256:823b65d8706e32ad2df51ed89496147a42a2a6e01c13cfb6ffb8b1e92bc910bb", size = 17127 },
+    { url = "https://files.pythonhosted.org/packages/a7/88/a940e11827ea1c136a34eca862486178294ae841164475b9ab216b80eb8e/MarkupSafe-2.1.5-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:c8b29db45f8fe46ad280a7294f5c3ec36dbac9491f2d1c17345be8e69cc5928f", size = 13982 },
+    { url = "https://files.pythonhosted.org/packages/cb/06/0d28bd178db529c5ac762a625c335a9168a7a23f280b4db9c95e97046145/MarkupSafe-2.1.5-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ec6a563cff360b50eed26f13adc43e61bc0c04d94b8be985e6fb24b81f6dcfdf", size = 26335 },
+    { url = "https://files.pythonhosted.org/packages/4a/1d/c4f5016f87ced614eacc7d5fb85b25bcc0ff53e8f058d069fc8cbfdc3c7a/MarkupSafe-2.1.5-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a549b9c31bec33820e885335b451286e2969a2d9e24879f83fe904a5ce59d70a", size = 25557 },
+    { url = "https://files.pythonhosted.org/packages/b3/fb/c18b8c9fbe69e347fdbf782c6478f1bc77f19a830588daa224236678339b/MarkupSafe-2.1.5-cp37-cp37m-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:4f11aa001c540f62c6166c7726f71f7573b52c68c31f014c25cc7901deea0b52", size = 25245 },
+    { url = "https://files.pythonhosted.org/packages/2f/69/30d29adcf9d1d931c75001dd85001adad7374381c9c2086154d9f6445be6/MarkupSafe-2.1.5-cp37-cp37m-musllinux_1_1_aarch64.whl", hash = "sha256:7b2e5a267c855eea6b4283940daa6e88a285f5f2a67f2220203786dfa59b37e9", size = 31013 },
+    { url = "https://files.pythonhosted.org/packages/3a/03/63498d05bd54278b6ca340099e5b52ffb9cdf2ee4f2d9b98246337e21689/MarkupSafe-2.1.5-cp37-cp37m-musllinux_1_1_i686.whl", hash = "sha256:2d2d793e36e230fd32babe143b04cec8a8b3eb8a3122d2aceb4a371e6b09b8df", size = 30178 },
+    { url = "https://files.pythonhosted.org/packages/68/79/11b4fe15124692f8673b603433e47abca199a08ecd2a4851bfbdc97dc62d/MarkupSafe-2.1.5-cp37-cp37m-musllinux_1_1_x86_64.whl", hash = "sha256:ce409136744f6521e39fd8e2a24c53fa18ad67aa5bc7c2cf83645cce5b5c4e50", size = 30429 },
+    { url = "https://files.pythonhosted.org/packages/ed/88/408bdbf292eb86f03201c17489acafae8358ba4e120d92358308c15cea7c/MarkupSafe-2.1.5-cp37-cp37m-win32.whl", hash = "sha256:4096e9de5c6fdf43fb4f04c26fb114f61ef0bf2e5604b6ee3019d51b69e8c371", size = 16633 },
+    { url = "https://files.pythonhosted.org/packages/6c/4c/3577a52eea1880538c435176bc85e5b3379b7ab442327ccd82118550758f/MarkupSafe-2.1.5-cp37-cp37m-win_amd64.whl", hash = "sha256:4275d846e41ecefa46e2015117a9f491e57a71ddd59bbead77e904dc02b1bed2", size = 17215 },
+    { url = "https://files.pythonhosted.org/packages/f8/ff/2c942a82c35a49df5de3a630ce0a8456ac2969691b230e530ac12314364c/MarkupSafe-2.1.5-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:656f7526c69fac7f600bd1f400991cc282b417d17539a1b228617081106feb4a", size = 18192 },
+    { url = "https://files.pythonhosted.org/packages/4f/14/6f294b9c4f969d0c801a4615e221c1e084722ea6114ab2114189c5b8cbe0/MarkupSafe-2.1.5-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:97cafb1f3cbcd3fd2b6fbfb99ae11cdb14deea0736fc2b0952ee177f2b813a46", size = 14072 },
+    { url = "https://files.pythonhosted.org/packages/81/d4/fd74714ed30a1dedd0b82427c02fa4deec64f173831ec716da11c51a50aa/MarkupSafe-2.1.5-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1f3fbcb7ef1f16e48246f704ab79d79da8a46891e2da03f8783a5b6fa41a9532", size = 26928 },
+    { url = "https://files.pythonhosted.org/packages/c7/bd/50319665ce81bb10e90d1cf76f9e1aa269ea6f7fa30ab4521f14d122a3df/MarkupSafe-2.1.5-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fa9db3f79de01457b03d4f01b34cf91bc0048eb2c3846ff26f66687c2f6d16ab", size = 26106 },
+    { url = "https://files.pythonhosted.org/packages/4c/6f/f2b0f675635b05f6afd5ea03c094557bdb8622fa8e673387444fe8d8e787/MarkupSafe-2.1.5-cp38-cp38-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ffee1f21e5ef0d712f9033568f8344d5da8cc2869dbd08d87c84656e6a2d2f68", size = 25781 },
+    { url = "https://files.pythonhosted.org/packages/51/e0/393467cf899b34a9d3678e78961c2c8cdf49fb902a959ba54ece01273fb1/MarkupSafe-2.1.5-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:5dedb4db619ba5a2787a94d877bc8ffc0566f92a01c0ef214865e54ecc9ee5e0", size = 30518 },
+    { url = "https://files.pythonhosted.org/packages/f6/02/5437e2ad33047290dafced9df741d9efc3e716b75583bbd73a9984f1b6f7/MarkupSafe-2.1.5-cp38-cp38-musllinux_1_1_i686.whl", hash = "sha256:30b600cf0a7ac9234b2638fbc0fb6158ba5bdcdf46aeb631ead21248b9affbc4", size = 29669 },
+    { url = "https://files.pythonhosted.org/packages/0e/7d/968284145ffd9d726183ed6237c77938c021abacde4e073020f920e060b2/MarkupSafe-2.1.5-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:8dd717634f5a044f860435c1d8c16a270ddf0ef8588d4887037c5028b859b0c3", size = 29933 },
+    { url = "https://files.pythonhosted.org/packages/bf/f3/ecb00fc8ab02b7beae8699f34db9357ae49d9f21d4d3de6f305f34fa949e/MarkupSafe-2.1.5-cp38-cp38-win32.whl", hash = "sha256:daa4ee5a243f0f20d528d939d06670a298dd39b1ad5f8a72a4275124a7819eff", size = 16656 },
+    { url = "https://files.pythonhosted.org/packages/92/21/357205f03514a49b293e214ac39de01fadd0970a6e05e4bf1ddd0ffd0881/MarkupSafe-2.1.5-cp38-cp38-win_amd64.whl", hash = "sha256:619bc166c4f2de5caa5a633b8b7326fbe98e0ccbfacabd87268a2b15ff73a029", size = 17206 },
+    { url = "https://files.pythonhosted.org/packages/0f/31/780bb297db036ba7b7bbede5e1d7f1e14d704ad4beb3ce53fb495d22bc62/MarkupSafe-2.1.5-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:7a68b554d356a91cce1236aa7682dc01df0edba8d043fd1ce607c49dd3c1edcf", size = 18193 },
+    { url = "https://files.pythonhosted.org/packages/6c/77/d77701bbef72892affe060cdacb7a2ed7fd68dae3b477a8642f15ad3b132/MarkupSafe-2.1.5-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:db0b55e0f3cc0be60c1f19efdde9a637c32740486004f20d1cff53c3c0ece4d2", size = 14073 },
+    { url = "https://files.pythonhosted.org/packages/d9/a7/1e558b4f78454c8a3a0199292d96159eb4d091f983bc35ef258314fe7269/MarkupSafe-2.1.5-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3e53af139f8579a6d5f7b76549125f0d94d7e630761a2111bc431fd820e163b8", size = 26486 },
+    { url = "https://files.pythonhosted.org/packages/5f/5a/360da85076688755ea0cceb92472923086993e86b5613bbae9fbc14136b0/MarkupSafe-2.1.5-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:17b950fccb810b3293638215058e432159d2b71005c74371d784862b7e4683f3", size = 25685 },
+    { url = "https://files.pythonhosted.org/packages/6a/18/ae5a258e3401f9b8312f92b028c54d7026a97ec3ab20bfaddbdfa7d8cce8/MarkupSafe-2.1.5-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:4c31f53cdae6ecfa91a77820e8b151dba54ab528ba65dfd235c80b086d68a465", size = 25338 },
+    { url = "https://files.pythonhosted.org/packages/0b/cc/48206bd61c5b9d0129f4d75243b156929b04c94c09041321456fd06a876d/MarkupSafe-2.1.5-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:bff1b4290a66b490a2f4719358c0cdcd9bafb6b8f061e45c7a2460866bf50c2e", size = 30439 },
+    { url = "https://files.pythonhosted.org/packages/d1/06/a41c112ab9ffdeeb5f77bc3e331fdadf97fa65e52e44ba31880f4e7f983c/MarkupSafe-2.1.5-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:bc1667f8b83f48511b94671e0e441401371dfd0f0a795c7daa4a3cd1dde55bea", size = 29531 },
+    { url = "https://files.pythonhosted.org/packages/02/8c/ab9a463301a50dab04d5472e998acbd4080597abc048166ded5c7aa768c8/MarkupSafe-2.1.5-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:5049256f536511ee3f7e1b3f87d1d1209d327e818e6ae1365e8653d7e3abb6a6", size = 29823 },
+    { url = "https://files.pythonhosted.org/packages/bc/29/9bc18da763496b055d8e98ce476c8e718dcfd78157e17f555ce6dd7d0895/MarkupSafe-2.1.5-cp39-cp39-win32.whl", hash = "sha256:00e046b6dd71aa03a41079792f8473dc494d564611a8f89bbbd7cb93295ebdcf", size = 16658 },
+    { url = "https://files.pythonhosted.org/packages/f6/f8/4da07de16f10551ca1f640c92b5f316f9394088b183c6a57183df6de5ae4/MarkupSafe-2.1.5-cp39-cp39-win_amd64.whl", hash = "sha256:fa173ec60341d6bb97a89f5ea19c85c5643c1e7dedebc22f5181eb73573142c5", size = 17211 },
+]
+
+[[package]]
+name = "nodeenv"
+version = "1.9.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/43/16/fc88b08840de0e0a72a2f9d8c6bae36be573e475a6326ae854bcc549fc45/nodeenv-1.9.1.tar.gz", hash = "sha256:6ec12890a2dab7946721edbfbcd91f3319c6ccc9aec47be7c7e6b7011ee6645f", size = 47437 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/d2/1d/1b658dbd2b9fa9c4c9f32accbfc0205d532c8c6194dc0f2a4c0428e7128a/nodeenv-1.9.1-py2.py3-none-any.whl", hash = "sha256:ba11c9782d29c27c70ffbdda2d7415098754709be8a7056d79a737cd901155c9", size = 22314 },
+]
+
+[[package]]
+name = "numpy"
+version = "1.21.6"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/45/b7/de7b8e67f2232c26af57c205aaad29fe17754f793404f59c8a730c7a191a/numpy-1.21.6.zip", hash = "sha256:ecb55251139706669fdec2ff073c98ef8e9a84473e51e716211b41aa0f18e656", size = 10274544 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/ff/c6/05ae3c7f75b596e1bb3d78131c331eada9376a03d1af9801bd40e4675023/numpy-1.21.6-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:8737609c3bbdd48e380d463134a35ffad3b22dc56295eff6f79fd85bd0eeeb25", size = 27203123 },
+    { url = "https://files.pythonhosted.org/packages/4a/72/a3379f83172f1431d7949138373e3a24beed68184c9362dab1b4d465be26/numpy-1.21.6-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:fdffbfb6832cd0b300995a2b08b8f6fa9f6e856d562800fea9182316d99c4e8e", size = 16965603 },
+    { url = "https://files.pythonhosted.org/packages/26/e7/4a6f579af8186372b03e8480e47df309520d91cfead8759b64dd5ac62688/numpy-1.21.6-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:3820724272f9913b597ccd13a467cc492a0da6b05df26ea09e78b171a0bb9da6", size = 12364595 },
+    { url = "https://files.pythonhosted.org/packages/57/ba/d8cbdfd507b541bb247beff24d9d7304ac8ffc379cf585701187d45d4512/numpy-1.21.6-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f17e562de9edf691a42ddb1eb4a5541c20dd3f9e65b09ded2beb0799c0cf29bb", size = 13019278 },
+    { url = "https://files.pythonhosted.org/packages/6f/7b/036000a55680e6c7eb81502b0aa27ce0ed65d4d8805613909967d9f8baf6/numpy-1.21.6-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5f30427731561ce75d7048ac254dbe47a2ba576229250fb60f0fb74db96501a1", size = 15906004 },
+    { url = "https://files.pythonhosted.org/packages/b0/77/ff8bbe56ff6cbbdbdb8a641c67cee61e29b2e8bfbb18732c2e1d2961fe4d/numpy-1.21.6-cp310-cp310-win32.whl", hash = "sha256:d4bf4d43077db55589ffc9009c0ba0a94fa4908b9586d6ccce2e0b164c86303c", size = 11706182 },
+    { url = "https://files.pythonhosted.org/packages/ec/03/93702ca9c4bd61791e46c80ff1f24943febb2317484cf7e8207688bbbd95/numpy-1.21.6-cp310-cp310-win_amd64.whl", hash = "sha256:d136337ae3cc69aa5e447e78d8e1514be8c3ec9b54264e680cf0b4bd9011574f", size = 14008913 },
+    { url = "https://files.pythonhosted.org/packages/32/dd/43d8b2b2ebf424f6555271a4c9f5b50dc3cc0aafa66c72b4d36863f71358/numpy-1.21.6-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:6aaf96c7f8cebc220cdfc03f1d5a31952f027dda050e5a703a0d1c396075e3e7", size = 16894122 },
+    { url = "https://files.pythonhosted.org/packages/06/78/b184f13f5461812a17a90b380d70a93fa3532460f0af9d72b0d93d8bc4ff/numpy-1.21.6-cp37-cp37m-manylinux_2_12_i686.manylinux2010_i686.whl", hash = "sha256:67c261d6c0a9981820c3a149d255a76918278a6b03b6a036800359aba1256d46", size = 13670697 },
+    { url = "https://files.pythonhosted.org/packages/6d/ad/ff3b21ebfe79a4d25b4a4f8e5cf9fd44a204adb6b33c09010f566f51027a/numpy-1.21.6-cp37-cp37m-manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:a6be4cb0ef3b8c9250c19cc122267263093eee7edd4e3fa75395dfda8c17a8e2", size = 15702369 },
+    { url = "https://files.pythonhosted.org/packages/b7/0d/86662f93102e42545cdf031da4fddf0ace9030ec67478932a628afc5973b/numpy-1.21.6-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7c4068a8c44014b2d55f3c3f574c376b2494ca9cc73d2f1bd692382b6dffe3db", size = 12974170 },
+    { url = "https://files.pythonhosted.org/packages/cd/eb/f6f3258e7b0e0cc5c327778312bf4ee4978c8514aa28e97119ee206f6e60/numpy-1.21.6-cp37-cp37m-win32.whl", hash = "sha256:7c7e5fa88d9ff656e067876e4736379cc962d185d5cd808014a8a928d529ef4e", size = 11680505 },
+    { url = "https://files.pythonhosted.org/packages/97/9f/da37cc4a188a1d5d203d65ab28d6504e17594b5342e0c1dc5610ee6f4535/numpy-1.21.6-cp37-cp37m-win_amd64.whl", hash = "sha256:bcb238c9c96c00d3085b264e5c1a1207672577b93fa666c3b14a45240b14123a", size = 13969236 },
+    { url = "https://files.pythonhosted.org/packages/b5/e2/b2df1f664d644e690b40179fc0a07c163c6decf986c7adee8a85a094e8ce/numpy-1.21.6-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:82691fda7c3f77c90e62da69ae60b5ac08e87e775b09813559f8901a88266552", size = 27127297 },
+    { url = "https://files.pythonhosted.org/packages/5b/d4/be63d2bed7d10f443dee42469623326b6bc51c9e5cd096ebb7227bca456f/numpy-1.21.6-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:643843bcc1c50526b3a71cd2ee561cf0d8773f062c8cbaf9ffac9fdf573f83ab", size = 16926225 },
+    { url = "https://files.pythonhosted.org/packages/0d/21/036363516c06737135ee58741e9c0af4899348ce3c5f5e04379240edd090/numpy-1.21.6-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:357768c2e4451ac241465157a3e929b265dfac85d9214074985b1786244f2ef3", size = 12329630 },
+    { url = "https://files.pythonhosted.org/packages/6a/52/a1dcf14b8e81d49c14112663290ee2ed545bd04988170138284a613bd926/numpy-1.21.6-cp38-cp38-manylinux_2_12_i686.manylinux2010_i686.whl", hash = "sha256:9f411b2c3f3d76bba0865b35a425157c5dcf54937f82bbeb3d3c180789dd66a6", size = 13719672 },
+    { url = "https://files.pythonhosted.org/packages/d5/43/e88bb1fb7d040ae8e0e06e749341b13f57701aab11fe9d71c99af6202c5c/numpy-1.21.6-cp38-cp38-manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:4aa48afdce4660b0076a00d80afa54e8a97cd49f457d68a4342d188a09451c1a", size = 15749422 },
+    { url = "https://files.pythonhosted.org/packages/86/c7/3f68d0a8dcc9458879c614707e6ffaf64a108664cfbba9702d3ba7ca4c82/numpy-1.21.6-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d6a96eef20f639e6a97d23e57dd0c1b1069a7b4fd7027482a4c5c451cd7732f4", size = 13026205 },
+    { url = "https://files.pythonhosted.org/packages/6f/47/453023bd298f8b0be092d8a8bdd4b21f87a8c639ecb724a94cd75e23d216/numpy-1.21.6-cp38-cp38-win32.whl", hash = "sha256:5c3c8def4230e1b959671eb959083661b4a0d2e9af93ee339c7dada6759a9470", size = 11707551 },
+    { url = "https://files.pythonhosted.org/packages/48/5f/db4550e1c68206814a577ebd92c0dd082f3628fd7fc96725d44a521b0c92/numpy-1.21.6-cp38-cp38-win_amd64.whl", hash = "sha256:bf2ec4b75d0e9356edea834d1de42b31fe11f726a81dfb2c2112bc1eaa508fcf", size = 14009785 },
+    { url = "https://files.pythonhosted.org/packages/83/eb/a6a0d7fc8e718776c5c710692ea027607104710cba813c4b869182179334/numpy-1.21.6-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:4391bd07606be175aafd267ef9bea87cf1b8210c787666ce82073b05f202add1", size = 27202394 },
+    { url = "https://files.pythonhosted.org/packages/4c/62/07402945bd5d5cf515a5f0cbc7263abf02ec0ddf3b19fbdc4af7537cd4d0/numpy-1.21.6-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:67f21981ba2f9d7ba9ade60c9e8cbaa8cf8e9ae51673934480e45cf55e953673", size = 16965163 },
+    { url = "https://files.pythonhosted.org/packages/44/56/041e886b4a8da813b7ec297c270fb3582d2ae8b7f33e106eb5c7a5e9184c/numpy-1.21.6-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:ee5ec40fdd06d62fe5d4084bef4fd50fd4bb6bfd2bf519365f569dc470163ab0", size = 12364236 },
+    { url = "https://files.pythonhosted.org/packages/61/f4/f01a8989e53a437ad660ab86c91514bec3d5067393e4a844b259f5a103de/numpy-1.21.6-cp39-cp39-manylinux_2_12_i686.manylinux2010_i686.whl", hash = "sha256:1dbe1c91269f880e364526649a52eff93ac30035507ae980d2fed33aaee633ac", size = 13721583 },
+    { url = "https://files.pythonhosted.org/packages/e7/f2/0bdcf2c40ef144cbbc9e0947eea831a145a98b0e5f8438fc09cf7fda0b35/numpy-1.21.6-cp39-cp39-manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:d9caa9d5e682102453d96a0ee10c7241b72859b01a941a397fd965f23b3e016b", size = 15734025 },
+    { url = "https://files.pythonhosted.org/packages/76/7f/830cf169eede1b855538f962e3a70c31755db6423652695b813ed04ff54e/numpy-1.21.6-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:58459d3bad03343ac4b1b42ed14d571b8743dc80ccbf27444f266729df1d6f5b", size = 13020515 },
+    { url = "https://files.pythonhosted.org/packages/1b/b5/7178d5a22427a9195ac69d6ec150415734f7a7a19d1142f82b89ead1dac4/numpy-1.21.6-cp39-cp39-win32.whl", hash = "sha256:7f5ae4f304257569ef3b948810816bc87c9146e8c446053539947eedeaa32786", size = 11706846 },
+    { url = "https://files.pythonhosted.org/packages/4d/04/bcd62448f2e772bc90a73ba21bacaa19817ae9905ae639969462862bd071/numpy-1.21.6-cp39-cp39-win_amd64.whl", hash = "sha256:e31f0bb5928b793169b87e3d1e070f2342b22d5245c755e2b81caa29756246c3", size = 14008965 },
+    { url = "https://files.pythonhosted.org/packages/2e/5a/6f3e280a10de48395053a559bfcb3b2221b74b57d062c1d6307fc965f549/numpy-1.21.6-pp37-pypy37_pp73-manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:dd1c8f6bd65d07d3810b90d02eba7997e32abbdf1277a481d698969e921a3be0", size = 15159527 },
+]
+
+[[package]]
+name = "packaging"
+version = "24.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/ee/b5/b43a27ac7472e1818c4bafd44430e69605baefe1f34440593e0332ec8b4d/packaging-24.0.tar.gz", hash = "sha256:eb82c5e3e56209074766e6885bb04b8c38a0c015d0a30036ebe7ece34c9989e9", size = 147882 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/49/df/1fceb2f8900f8639e278b056416d49134fb8d84c5942ffaa01ad34782422/packaging-24.0-py3-none-any.whl", hash = "sha256:2ddfb553fdf02fb784c234c7ba6ccc288296ceabec964ad2eae3777778130bc5", size = 53488 },
+]
+
+[[package]]
+name = "platformdirs"
+version = "4.0.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "typing-extensions", marker = "python_full_version < '3.8'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/31/28/e40d24d2e2eb23135f8533ad33d582359c7825623b1e022f9d460def7c05/platformdirs-4.0.0.tar.gz", hash = "sha256:cb633b2bcf10c51af60beb0ab06d2f1d69064b43abf4c185ca6b28865f3f9731", size = 19914 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/31/16/70be3b725073035aa5fc3229321d06e22e73e3e09f6af78dcfdf16c7636c/platformdirs-4.0.0-py3-none-any.whl", hash = "sha256:118c954d7e949b35437270383a3f2531e99dd93cf7ce4dc8340d3356d30f173b", size = 17562 },
+]
+
+[[package]]
+name = "pluggy"
+version = "1.2.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "importlib-metadata", marker = "python_full_version < '3.8'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/8a/42/8f2833655a29c4e9cb52ee8a2be04ceac61bcff4a680fb338cbd3d1e322d/pluggy-1.2.0.tar.gz", hash = "sha256:d12f0c4b579b15f5e054301bb226ee85eeeba08ffec228092f8defbaa3a4c4b3", size = 61613 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/51/32/4a79112b8b87b21450b066e102d6608907f4c885ed7b04c3fdb085d4d6ae/pluggy-1.2.0-py3-none-any.whl", hash = "sha256:c2fd55a7d7a3863cba1a013e4e2414658b1d07b6bc57b3919e0c63c9abb99849", size = 17695 },
+]
+
+[[package]]
+name = "pre-commit"
+version = "2.21.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "cfgv" },
+    { name = "identify" },
+    { name = "importlib-metadata", marker = "python_full_version < '3.8'" },
+    { name = "nodeenv" },
+    { name = "pyyaml" },
+    { name = "virtualenv" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/6b/00/1637ae945c6e10838ef5c41965f1c864e59301811bb203e979f335608e7c/pre_commit-2.21.0.tar.gz", hash = "sha256:31ef31af7e474a8d8995027fefdfcf509b5c913ff31f2015b4ec4beb26a6f658", size = 174966 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/a6/6b/6cfe3a8b351b54f4b6c6d2ad4286804e3367f628dce379c603d3b96635f4/pre_commit-2.21.0-py2.py3-none-any.whl", hash = "sha256:e2f91727039fc39a92f58a588a25b87f936de6567eed4f0e673e0507edc75bad", size = 201938 },
+]
+
+[[package]]
+name = "pyarrow"
+version = "12.0.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "numpy" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/c5/68/d3410e975bebbf5be00c1238d0418345d8ec5d88b7a6c102211a1c967edd/pyarrow-12.0.1.tar.gz", hash = "sha256:cce317fc96e5b71107bf1f9f184d5e54e2bd14bbf3f9a3d62819961f0af86fec", size = 1015259 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/13/2f/a42dbdf34528c70bbd5736a968631e3c8c2f911aea89f9c49f6f834e83b5/pyarrow-12.0.1-cp310-cp310-macosx_10_14_x86_64.whl", hash = "sha256:6d288029a94a9bb5407ceebdd7110ba398a00412c5b0155ee9813a40d246c5df", size = 24731979 },
+    { url = "https://files.pythonhosted.org/packages/64/05/76bcbea6903957c6467f99fcc6aaf07ac5ea675c02e75881719949801335/pyarrow-12.0.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:345e1828efdbd9aa4d4de7d5676778aba384a2c3add896d995b23d368e60e5af", size = 22646475 },
+    { url = "https://files.pythonhosted.org/packages/60/94/e56483c49ae2acee47af880ab4e0af7749811a0142a584d45543957ee1b3/pyarrow-12.0.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8d6009fdf8986332b2169314da482baed47ac053311c8934ac6651e614deacd6", size = 36386496 },
+    { url = "https://files.pythonhosted.org/packages/0d/c8/886acfcce7cb2f7552f538d2b6deafd4841f3de42902943db15f1b42313d/pyarrow-12.0.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2d3c4cbbf81e6dd23fe921bc91dc4619ea3b79bc58ef10bce0f49bdafb103daf", size = 38938438 },
+    { url = "https://files.pythonhosted.org/packages/13/68/1bc83fe2de87e2b785c503ef28293b56a4405d632f2b443cf00265b12d90/pyarrow-12.0.1-cp310-cp310-win_amd64.whl", hash = "sha256:cdacf515ec276709ac8042c7d9bd5be83b4f5f39c6c037a17a60d7ebfd92c890", size = 21465027 },
+    { url = "https://files.pythonhosted.org/packages/90/1e/fb0177d214a77198083156d750358c0a3ff696c96b329f443ad5513d25b6/pyarrow-12.0.1-cp311-cp311-macosx_10_14_x86_64.whl", hash = "sha256:749be7fd2ff260683f9cc739cb862fb11be376de965a2a8ccbf2693b098db6c7", size = 24687651 },
+    { url = "https://files.pythonhosted.org/packages/a7/ca/a34c5dd3393644865b82ac5df66e52311fd4ae2fc073f62b68b8538a0da4/pyarrow-12.0.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:6895b5fb74289d055c43db3af0de6e16b07586c45763cb5e558d38b86a91e3a7", size = 22601949 },
+    { url = "https://files.pythonhosted.org/packages/6a/ba/571de5dc75831b9a0f9e8d23823c1286b5c940588d4d8c87aab535779d53/pyarrow-12.0.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1887bdae17ec3b4c046fcf19951e71b6a619f39fa674f9881216173566c8f718", size = 36384544 },
+    { url = "https://files.pythonhosted.org/packages/c5/52/19832487e6834164c523386a1b047dd5539fcbb876196b6f5619dfdab465/pyarrow-12.0.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e2c9cb8eeabbadf5fcfc3d1ddea616c7ce893db2ce4dcef0ac13b099ad7ca082", size = 38930218 },
+    { url = "https://files.pythonhosted.org/packages/53/98/823deb5d3ff75386d5ab19b90c25b4a3aca01299bb53f217f2cfa954329c/pyarrow-12.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:ce4aebdf412bd0eeb800d8e47db854f9f9f7e2f5a0220440acf219ddfddd4f63", size = 21453804 },
+    { url = "https://files.pythonhosted.org/packages/8f/56/10fab8ea743b9bfd954d8648e715e1a947d7e131858d9670f83770626059/pyarrow-12.0.1-cp37-cp37m-macosx_10_14_x86_64.whl", hash = "sha256:e0d8730c7f6e893f6db5d5b86eda42c0a130842d101992b581e2138e4d5663d3", size = 24707735 },
+    { url = "https://files.pythonhosted.org/packages/cf/e2/94791e4cbb8cb16d5c99d016003746fab9d97f127342ba6b817bf639c767/pyarrow-12.0.1-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:43364daec02f69fec89d2315f7fbfbeec956e0d991cbbef471681bd77875c40f", size = 36561574 },
+    { url = "https://files.pythonhosted.org/packages/00/bd/4c03789f723337723670e8cf8935b18e170e7af3813ec38e71d2cfb0bd93/pyarrow-12.0.1-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:051f9f5ccf585f12d7de836e50965b3c235542cc896959320d9776ab93f3b33d", size = 39110664 },
+    { url = "https://files.pythonhosted.org/packages/0b/82/7f70296eb5167bc3bcee96a1460315af109affd7fef43e750c8ee4aac17c/pyarrow-12.0.1-cp37-cp37m-win_amd64.whl", hash = "sha256:be2757e9275875d2a9c6e6052ac7957fbbfc7bc7370e4a036a9b893e96fedaba", size = 21471413 },
+    { url = "https://files.pythonhosted.org/packages/75/a2/87fe24ab2c6efc6ad2335a2fc6bc33363fc70f67f18a3c18c494a4783aa2/pyarrow-12.0.1-cp38-cp38-macosx_10_14_x86_64.whl", hash = "sha256:cf812306d66f40f69e684300f7af5111c11f6e0d89d6b733e05a3de44961529d", size = 24719478 },
+    { url = "https://files.pythonhosted.org/packages/01/7c/427358d04bfcb217a165911bfada09e2d1e74fed04e40eb02c0c317ed2c4/pyarrow-12.0.1-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:459a1c0ed2d68671188b2118c63bac91eaef6fc150c77ddd8a583e3c795737bf", size = 22638986 },
+    { url = "https://files.pythonhosted.org/packages/25/72/9afe7e2b61482ddc361c796857c19f69b2035ae20deebe0c1a54cb602b21/pyarrow-12.0.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:85e705e33eaf666bbe508a16fd5ba27ca061e177916b7a317ba5a51bee43384c", size = 36490338 },
+    { url = "https://files.pythonhosted.org/packages/e1/91/676b6ef5181fd0229ec35477eb94ff55fc5114ebab7a4669db311ddc9385/pyarrow-12.0.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9120c3eb2b1f6f516a3b7a9714ed860882d9ef98c4b17edcdc91d95b7528db60", size = 39043549 },
+    { url = "https://files.pythonhosted.org/packages/af/cd/9674a609185bb9197b2dc25e8e61e1f4a2531f0754f43e4d3790f885707e/pyarrow-12.0.1-cp38-cp38-win_amd64.whl", hash = "sha256:c780f4dc40460015d80fcd6a6140de80b615349ed68ef9adb653fe351778c9b3", size = 21517590 },
+    { url = "https://files.pythonhosted.org/packages/25/33/8fa80189ea3ea7ac0b35b33e715de0466a0ec5064abb07a5b7ab5fe4f6fe/pyarrow-12.0.1-cp39-cp39-macosx_10_14_x86_64.whl", hash = "sha256:a3c63124fc26bf5f95f508f5d04e1ece8cc23a8b0af2a1e6ab2b1ec3fdc91b24", size = 24770513 },
+    { url = "https://files.pythonhosted.org/packages/a7/fd/a1488faf625a86b2ebf83bb977e48d9514785edfe438d4dbccf6e527bcc8/pyarrow-12.0.1-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:b13329f79fa4472324f8d32dc1b1216616d09bd1e77cfb13104dec5463632c36", size = 22676028 },
+    { url = "https://files.pythonhosted.org/packages/8b/14/dbda2f416906090824e5b58134ebef504065798bbcc98c929ce712be80ed/pyarrow-12.0.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bb656150d3d12ec1396f6dde542db1675a95c0cc8366d507347b0beed96e87ca", size = 36430612 },
+    { url = "https://files.pythonhosted.org/packages/54/a2/5976df95323c4ca2b7baba31cb7a2a61a17461706043239d38a8e9dc281e/pyarrow-12.0.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6251e38470da97a5b2e00de5c6a049149f7b2bd62f12fa5dbb9ac674119ba71a", size = 38983701 },
+    { url = "https://files.pythonhosted.org/packages/dc/45/31441c988329afed625a791a7d78f1cf2fcb40dcc86a1d61e081287516a8/pyarrow-12.0.1-cp39-cp39-win_amd64.whl", hash = "sha256:3de26da901216149ce086920547dfff5cd22818c9eab67ebc41e863a5883bac7", size = 21511440 },
+]
+
+[[package]]
+name = "pycparser"
+version = "2.21"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/5e/0b/95d387f5f4433cb0f53ff7ad859bd2c6051051cebbb564f139a999ab46de/pycparser-2.21.tar.gz", hash = "sha256:e644fdec12f7872f86c58ff790da456218b10f863970249516d60a5eaca77206", size = 170877 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/62/d5/5f610ebe421e85889f2e55e33b7f9a6795bd982198517d912eb1c76e1a53/pycparser-2.21-py2.py3-none-any.whl", hash = "sha256:8ee45429555515e1f6b185e78100aea234072576aa43ab53aefcae078162fca9", size = 118697 },
+]
+
+[[package]]
+name = "pytest"
+version = "7.4.4"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "colorama", marker = "sys_platform == 'win32'" },
+    { name = "exceptiongroup", marker = "python_full_version < '3.11'" },
+    { name = "importlib-metadata", marker = "python_full_version < '3.8'" },
+    { name = "iniconfig" },
+    { name = "packaging" },
+    { name = "pluggy" },
+    { name = "tomli", marker = "python_full_version < '3.11'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/80/1f/9d8e98e4133ffb16c90f3b405c43e38d3abb715bb5d7a63a5a684f7e46a3/pytest-7.4.4.tar.gz", hash = "sha256:2cf0005922c6ace4a3e2ec8b4080eb0d9753fdc93107415332f50ce9e7994280", size = 1357116 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/51/ff/f6e8b8f39e08547faece4bd80f89d5a8de68a38b2d179cc1c4490ffa3286/pytest-7.4.4-py3-none-any.whl", hash = "sha256:b090cdf5ed60bf4c45261be03239c2c1c22df034fbffe691abe93cd80cea01d8", size = 325287 },
+]
+
+[[package]]
+name = "pyyaml"
+version = "6.0.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/cd/e5/af35f7ea75cf72f2cd079c95ee16797de7cd71f29ea7c68ae5ce7be1eda0/PyYAML-6.0.1.tar.gz", hash = "sha256:bfdf460b1736c775f2ba9f6a92bca30bc2095067b8a9d77876d1fad6cc3b4a43", size = 125201 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/96/06/4beb652c0fe16834032e54f0956443d4cc797fe645527acee59e7deaa0a2/PyYAML-6.0.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:d858aa552c999bc8a8d57426ed01e40bef403cd8ccdd0fc5f6f04a00414cac2a", size = 189447 },
+    { url = "https://files.pythonhosted.org/packages/5b/07/10033a403b23405a8fc48975444463d3d10a5c2736b7eb2550b07b367429/PyYAML-6.0.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:fd66fc5d0da6d9815ba2cebeb4205f95818ff4b79c3ebe268e75d961704af52f", size = 169264 },
+    { url = "https://files.pythonhosted.org/packages/f1/26/55e4f21db1f72eaef092015d9017c11510e7e6301c62a6cfee91295d13c6/PyYAML-6.0.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:69b023b2b4daa7548bcfbd4aa3da05b3a74b772db9e23b982788168117739938", size = 677003 },
+    { url = "https://files.pythonhosted.org/packages/ba/91/090818dfa62e85181f3ae23dd1e8b7ea7f09684864a900cab72d29c57346/PyYAML-6.0.1-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:81e0b275a9ecc9c0c0c07b4b90ba548307583c125f54d5b6946cfee6360c733d", size = 699070 },
+    { url = "https://files.pythonhosted.org/packages/29/61/bf33c6c85c55bc45a29eee3195848ff2d518d84735eb0e2d8cb42e0d285e/PyYAML-6.0.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ba336e390cd8e4d1739f42dfe9bb83a3cc2e80f567d8805e11b46f4a943f5515", size = 705525 },
+    { url = "https://files.pythonhosted.org/packages/07/91/45dfd0ef821a7f41d9d0136ea3608bb5b1653e42fd56a7970532cb5c003f/PyYAML-6.0.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:326c013efe8048858a6d312ddd31d56e468118ad4cdeda36c719bf5bb6192290", size = 707514 },
+    { url = "https://files.pythonhosted.org/packages/b6/a0/b6700da5d49e9fed49dc3243d3771b598dad07abb37cc32e524607f96adc/PyYAML-6.0.1-cp310-cp310-win32.whl", hash = "sha256:bd4af7373a854424dabd882decdc5579653d7868b8fb26dc7d0e99f823aa5924", size = 130488 },
+    { url = "https://files.pythonhosted.org/packages/24/97/9b59b43431f98d01806b288532da38099cc6f2fea0f3d712e21e269c0279/PyYAML-6.0.1-cp310-cp310-win_amd64.whl", hash = "sha256:fd1592b3fdf65fff2ad0004b5e363300ef59ced41c2e6b3a99d4089fa8c5435d", size = 145338 },
+    { url = "https://files.pythonhosted.org/packages/ec/0d/26fb23e8863e0aeaac0c64e03fd27367ad2ae3f3cccf3798ee98ce160368/PyYAML-6.0.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:6965a7bc3cf88e5a1c3bd2e0b5c22f8d677dc88a455344035f03399034eb3007", size = 187867 },
+    { url = "https://files.pythonhosted.org/packages/28/09/55f715ddbf95a054b764b547f617e22f1d5e45d83905660e9a088078fe67/PyYAML-6.0.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:f003ed9ad21d6a4713f0a9b5a7a0a79e08dd0f221aff4525a2be4c346ee60aab", size = 167530 },
+    { url = "https://files.pythonhosted.org/packages/5e/94/7d5ee059dfb92ca9e62f4057dcdec9ac08a9e42679644854dc01177f8145/PyYAML-6.0.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:42f8152b8dbc4fe7d96729ec2b99c7097d656dc1213a3229ca5383f973a5ed6d", size = 732244 },
+    { url = "https://files.pythonhosted.org/packages/06/92/e0224aa6ebf9dc54a06a4609da37da40bb08d126f5535d81bff6b417b2ae/PyYAML-6.0.1-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:062582fca9fabdd2c8b54a3ef1c978d786e0f6b3a1510e0ac93ef59e0ddae2bc", size = 752871 },
+    { url = "https://files.pythonhosted.org/packages/7b/5e/efd033ab7199a0b2044dab3b9f7a4f6670e6a52c089de572e928d2873b06/PyYAML-6.0.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d2b04aac4d386b172d5b9692e2d2da8de7bfb6c387fa4f801fbf6fb2e6ba4673", size = 757729 },
+    { url = "https://files.pythonhosted.org/packages/03/5c/c4671451b2f1d76ebe352c0945d4cd13500adb5d05f5a51ee296d80152f7/PyYAML-6.0.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:e7d73685e87afe9f3b36c799222440d6cf362062f78be1013661b00c5c6f678b", size = 748528 },
+    { url = "https://files.pythonhosted.org/packages/73/9c/766e78d1efc0d1fca637a6b62cea1b4510a7fb93617eb805223294fef681/PyYAML-6.0.1-cp311-cp311-win32.whl", hash = "sha256:1635fd110e8d85d55237ab316b5b011de701ea0f29d07611174a1b42f1444741", size = 130286 },
+    { url = "https://files.pythonhosted.org/packages/b3/34/65bb4b2d7908044963ebf614fe0fdb080773fc7030d7e39c8d3eddcd4257/PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34", size = 144699 },
+    { url = "https://files.pythonhosted.org/packages/bc/06/1b305bf6aa704343be85444c9d011f626c763abb40c0edc1cad13bfd7f86/PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28", size = 178692 },
+    { url = "https://files.pythonhosted.org/packages/84/02/404de95ced348b73dd84f70e15a41843d817ff8c1744516bf78358f2ffd2/PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9", size = 165622 },
+    { url = "https://files.pythonhosted.org/packages/c7/4c/4a2908632fc980da6d918b9de9c1d9d7d7e70b2672b1ad5166ed27841ef7/PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef", size = 696937 },
+    { url = "https://files.pythonhosted.org/packages/b4/33/720548182ffa8344418126017aa1d4ab4aeec9a2275f04ce3f3573d8ace8/PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0", size = 724969 },
+    { url = "https://files.pythonhosted.org/packages/4f/78/77b40157b6cb5f2d3d31a3d9b2efd1ba3505371f76730d267e8b32cf4b7f/PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4", size = 712604 },
+    { url = "https://files.pythonhosted.org/packages/2e/97/3e0e089ee85e840f4b15bfa00e4e63d84a3691ababbfea92d6f820ea6f21/PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54", size = 126098 },
+    { url = "https://files.pythonhosted.org/packages/2b/9f/fbade56564ad486809c27b322d0f7e6a89c01f6b4fe208402e90d4443a99/PyYAML-6.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:0d3304d8c0adc42be59c5f8a4d9e3d7379e6955ad754aa9d6ab7a398b59dd1df", size = 138675 },
+    { url = "https://files.pythonhosted.org/packages/c7/d1/02baa09d39b1bb1ebaf0d850d106d1bdcb47c91958557f471153c49dc03b/PyYAML-6.0.1-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:b1275ad35a5d18c62a7220633c913e1b42d44b46ee12554e5fd39c70a243d6a3", size = 189627 },
+    { url = "https://files.pythonhosted.org/packages/e5/31/ba812efa640a264dbefd258986a5e4e786230cb1ee4a9f54eb28ca01e14a/PyYAML-6.0.1-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:18aeb1bf9a78867dc38b259769503436b7c72f7a1f1f4c93ff9a17de54319b27", size = 658438 },
+    { url = "https://files.pythonhosted.org/packages/4d/f1/08f06159739254c8947899c9fc901241614195db15ba8802ff142237664c/PyYAML-6.0.1-cp37-cp37m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:596106435fa6ad000c2991a98fa58eeb8656ef2325d7e158344fb33864ed87e3", size = 680304 },
+    { url = "https://files.pythonhosted.org/packages/d7/8f/db62b0df635b9008fe90aa68424e99cee05e68b398740c8a666a98455589/PyYAML-6.0.1-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:baa90d3f661d43131ca170712d903e6295d1f7a0f595074f151c0aed377c9b9c", size = 670140 },
+    { url = "https://files.pythonhosted.org/packages/cc/5c/fcabd17918348c7db2eeeb0575705aaf3f7ab1657f6ce29b2e31737dd5d1/PyYAML-6.0.1-cp37-cp37m-win32.whl", hash = "sha256:9046c58c4395dff28dd494285c82ba00b546adfc7ef001486fbf0324bc174fba", size = 137577 },
+    { url = "https://files.pythonhosted.org/packages/1e/ae/964ccb88a938f20ece5754878f182cfbd846924930d02d29d06af8d4c69e/PyYAML-6.0.1-cp37-cp37m-win_amd64.whl", hash = "sha256:4fb147e7a67ef577a588a0e2c17b6db51dda102c71de36f8549b6816a96e1867", size = 153248 },
+    { url = "https://files.pythonhosted.org/packages/7f/5d/2779ea035ba1e533c32ed4a249b4e0448f583ba10830b21a3cddafe11a4e/PyYAML-6.0.1-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:1d4c7e777c441b20e32f52bd377e0c409713e8bb1386e1099c2415f26e479595", size = 191734 },
+    { url = "https://files.pythonhosted.org/packages/e1/a1/27bfac14b90adaaccf8c8289f441e9f76d94795ec1e7a8f134d9f2cb3d0b/PyYAML-6.0.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a0cd17c15d3bb3fa06978b4e8958dcdc6e0174ccea823003a106c7d4d7899ac5", size = 723767 },
+    { url = "https://files.pythonhosted.org/packages/c1/39/47ed4d65beec9ce07267b014be85ed9c204fa373515355d3efa62d19d892/PyYAML-6.0.1-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:28c119d996beec18c05208a8bd78cbe4007878c6dd15091efb73a30e90539696", size = 749067 },
+    { url = "https://files.pythonhosted.org/packages/c8/6b/6600ac24725c7388255b2f5add93f91e58a5d7efaf4af244fdbcc11a541b/PyYAML-6.0.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7e07cbde391ba96ab58e532ff4803f79c4129397514e1413a7dc761ccd755735", size = 736569 },
+    { url = "https://files.pythonhosted.org/packages/0d/46/62ae77677e532c0af6c81ddd6f3dbc16bdcc1208b077457354442d220bfb/PyYAML-6.0.1-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:49a183be227561de579b4a36efbb21b3eab9651dd81b1858589f796549873dd6", size = 787738 },
+    { url = "https://files.pythonhosted.org/packages/d6/6a/439d1a6f834b9a9db16332ce16c4a96dd0e3970b65fe08cbecd1711eeb77/PyYAML-6.0.1-cp38-cp38-win32.whl", hash = "sha256:184c5108a2aca3c5b3d3bf9395d50893a7ab82a38004c8f61c258d4428e80206", size = 139797 },
+    { url = "https://files.pythonhosted.org/packages/29/0f/9782fa5b10152abf033aec56a601177ead85ee03b57781f2d9fced09eefc/PyYAML-6.0.1-cp38-cp38-win_amd64.whl", hash = "sha256:1e2722cc9fbb45d9b87631ac70924c11d3a401b2d7f410cc0e3bbf249f2dca62", size = 157350 },
+    { url = "https://files.pythonhosted.org/packages/57/c5/5d09b66b41d549914802f482a2118d925d876dc2a35b2d127694c1345c34/PyYAML-6.0.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:9eb6caa9a297fc2c2fb8862bc5370d0303ddba53ba97e71f08023b6cd73d16a8", size = 197846 },
+    { url = "https://files.pythonhosted.org/packages/0e/88/21b2f16cb2123c1e9375f2c93486e35fdc86e63f02e274f0e99c589ef153/PyYAML-6.0.1-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:c8098ddcc2a85b61647b2590f825f3db38891662cfc2fc776415143f599bb859", size = 174396 },
+    { url = "https://files.pythonhosted.org/packages/ac/6c/967d91a8edf98d2b2b01d149bd9e51b8f9fb527c98d80ebb60c6b21d60c4/PyYAML-6.0.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5773183b6446b2c99bb77e77595dd486303b4faab2b086e7b17bc6bef28865f6", size = 731824 },
+    { url = "https://files.pythonhosted.org/packages/4a/4b/c71ef18ef83c82f99e6da8332910692af78ea32bd1d1d76c9787dfa36aea/PyYAML-6.0.1-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:b786eecbdf8499b9ca1d697215862083bd6d2a99965554781d0d8d1ad31e13a0", size = 754777 },
+    { url = "https://files.pythonhosted.org/packages/7d/39/472f2554a0f1e825bd7c5afc11c817cd7a2f3657460f7159f691fbb37c51/PyYAML-6.0.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bc1bf2925a1ecd43da378f4db9e4f799775d6367bdb94671027b73b393a7c42c", size = 738883 },
+    { url = "https://files.pythonhosted.org/packages/40/da/a175a35cf5583580e90ac3e2a3dbca90e43011593ae62ce63f79d7b28d92/PyYAML-6.0.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:04ac92ad1925b2cff1db0cfebffb6ffc43457495c9b3c39d3fcae417d7125dc5", size = 750294 },
+    { url = "https://files.pythonhosted.org/packages/24/62/7fcc372442ec8ea331da18c24b13710e010c5073ab851ef36bf9dacb283f/PyYAML-6.0.1-cp39-cp39-win32.whl", hash = "sha256:faca3bdcf85b2fc05d06ff3fbc1f83e1391b3e724afa3feba7d13eeab355484c", size = 136936 },
+    { url = "https://files.pythonhosted.org/packages/84/4d/82704d1ab9290b03da94e6425f5e87396b999fd7eb8e08f3a92c158402bf/PyYAML-6.0.1-cp39-cp39-win_amd64.whl", hash = "sha256:510c9deebc5c0225e8c96813043e62b680ba2f9c50a08d3724c7f28a747d1486", size = 152751 },
+]
+
+[[package]]
+name = "requests"
+version = "2.31.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "certifi" },
+    { name = "charset-normalizer" },
+    { name = "idna" },
+    { name = "urllib3" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/9d/be/10918a2eac4ae9f02f6cfe6414b7a155ccd8f7f9d4380d62fd5b955065c3/requests-2.31.0.tar.gz", hash = "sha256:942c5a758f98d790eaed1a29cb6eefc7ffb0d1cf7af05c3d2791656dbd6ad1e1", size = 110794 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/70/8e/0e2d847013cb52cd35b38c009bb167a1a26b2ce6cd6965bf26b47bc0bf44/requests-2.31.0-py3-none-any.whl", hash = "sha256:58cd2187c01e70e6e26505bca751777aa9f2ee0b7f4300988b709f44e013003f", size = 62574 },
+]
+
+[[package]]
+name = "setuptools"
+version = "68.0.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/dc/98/5f896af066c128669229ff1aa81553ac14cfb3e5e74b6b44594132b8540e/setuptools-68.0.0.tar.gz", hash = "sha256:baf1fdb41c6da4cd2eae722e135500da913332ab3f2f5c7d33af9b492acb5235", size = 2194111 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/c7/42/be1c7bbdd83e1bfb160c94b9cafd8e25efc7400346cf7ccdbdb452c467fa/setuptools-68.0.0-py3-none-any.whl", hash = "sha256:11e52c67415a381d10d6b462ced9cfb97066179f0e871399e006c4ab101fc85f", size = 804037 },
+]
+
+[[package]]
+name = "six"
+version = "1.17.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/94/e7/b2c673351809dca68a0e064b6af791aa332cf192da575fd474ed7d6f16a2/six-1.17.0.tar.gz", hash = "sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81", size = 34031 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050 },
+]
+
+[[package]]
+name = "tomli"
+version = "2.0.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/c0/3f/d7af728f075fb08564c5949a9c95e44352e23dee646869fa104a3b2060a3/tomli-2.0.1.tar.gz", hash = "sha256:de526c12914f0c550d15924c62d72abc48d6fe7364aa87328337a31007fe8a4f", size = 15164 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/97/75/10a9ebee3fd790d20926a90a2547f0bf78f371b2f13aa822c759680ca7b9/tomli-2.0.1-py3-none-any.whl", hash = "sha256:939de3e7a6161af0c887ef91b7d41a53e7c5a1ca976325f429cb46ea9bc30ecc", size = 12757 },
+]
+
+[[package]]
+name = "tqdm"
+version = "4.67.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "colorama", marker = "platform_system == 'Windows'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/a8/4b/29b4ef32e036bb34e4ab51796dd745cdba7ed47ad142a9f4a1eb8e0c744d/tqdm-4.67.1.tar.gz", hash = "sha256:f8aef9c52c08c13a65f30ea34f4e5aac3fd1a34959879d7e59e63027286627f2", size = 169737 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/d0/30/dc54f88dd4a2b5dc8a0279bdd7270e735851848b762aeb1c1184ed1f6b14/tqdm-4.67.1-py3-none-any.whl", hash = "sha256:26445eca388f82e72884e0d580d5464cd801a3ea01e63e5601bdff9ba6a48de2", size = 78540 },
+]
+
+[[package]]
+name = "triton"
+version = "3.3.1"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "setuptools" },
+]
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/8d/a9/549e51e9b1b2c9b854fd761a1d23df0ba2fbc60bd0c13b489ffa518cfcb7/triton-3.3.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b74db445b1c562844d3cfad6e9679c72e93fdfb1a90a24052b03bb5c49d1242e", size = 155600257 },
+    { url = "https://files.pythonhosted.org/packages/21/2f/3e56ea7b58f80ff68899b1dbe810ff257c9d177d288c6b0f55bf2fe4eb50/triton-3.3.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b31e3aa26f8cb3cc5bf4e187bf737cbacf17311e1112b781d4a059353dfd731b", size = 155689937 },
+    { url = "https://files.pythonhosted.org/packages/24/5f/950fb373bf9c01ad4eb5a8cd5eaf32cdf9e238c02f9293557a2129b9c4ac/triton-3.3.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9999e83aba21e1a78c1f36f21bce621b77bcaa530277a50484a7cb4a822f6e43", size = 155669138 },
+    { url = "https://files.pythonhosted.org/packages/74/1f/dfb531f90a2d367d914adfee771babbd3f1a5b26c3f5fbc458dee21daa78/triton-3.3.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b89d846b5a4198317fec27a5d3a609ea96b6d557ff44b56c23176546023c4240", size = 155673035 },
+    { url = "https://files.pythonhosted.org/packages/28/71/bd20ffcb7a64c753dc2463489a61bf69d531f308e390ad06390268c4ea04/triton-3.3.1-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a3198adb9d78b77818a5388bff89fa72ff36f9da0bc689db2f0a651a67ce6a42", size = 155735832 },
+    { url = "https://files.pythonhosted.org/packages/6d/81/ac4d50af22f594c4cb7c84fd2ad5ba1e0c03e2a83fe3483ddd79edcd7ec7/triton-3.3.1-cp39-cp39-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f6139aeb04a146b0b8e0fbbd89ad1e65861c57cfed881f21d62d3cb94a36bab7", size = 155596799 },
+]
+
+[[package]]
+name = "triton-viz"
+version = "2.0"
+source = { editable = "." }
+dependencies = [
+    { name = "anytree" },
+    { name = "cairocffi" },
+    { name = "flask" },
+    { name = "flask-cloudflared" },
+    { name = "pre-commit" },
+    { name = "pyarrow" },
+    { name = "pytest" },
+    { name = "setuptools" },
+    { name = "tqdm" },
+    { name = "triton" },
+    { name = "z3-solver" },
+]
+
+[package.metadata]
+requires-dist = [
+    { name = "anytree" },
+    { name = "cairocffi" },
+    { name = "flask" },
+    { name = "flask-cloudflared" },
+    { name = "pre-commit" },
+    { name = "pyarrow" },
+    { name = "pytest" },
+    { name = "setuptools" },
+    { name = "tqdm" },
+    { name = "triton" },
+    { name = "z3-solver" },
+]
+
+[[package]]
+name = "typing-extensions"
+version = "4.7.1"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/3c/8b/0111dd7d6c1478bf83baa1cab85c686426c7a6274119aceb2bd9d35395ad/typing_extensions-4.7.1.tar.gz", hash = "sha256:b75ddc264f0ba5615db7ba217daeb99701ad295353c45f9e95963337ceeeffb2", size = 72876 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/ec/6b/63cc3df74987c36fe26157ee12e09e8f9db4de771e0f3404263117e75b95/typing_extensions-4.7.1-py3-none-any.whl", hash = "sha256:440d5dd3af93b060174bf433bccd69b0babc3b15b1a8dca43789fd7f61514b36", size = 33232 },
+]
+
+[[package]]
+name = "urllib3"
+version = "2.0.7"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/af/47/b215df9f71b4fdba1025fc05a77db2ad243fa0926755a52c5e71659f4e3c/urllib3-2.0.7.tar.gz", hash = "sha256:c97dfde1f7bd43a71c8d2a58e369e9b2bf692d1334ea9f9cae55add7d0dd0f84", size = 282546 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/d2/b2/b157855192a68541a91ba7b2bbcb91f1b4faa51f8bae38d8005c034be524/urllib3-2.0.7-py3-none-any.whl", hash = "sha256:fdb6d215c776278489906c2f8916e6e7d4f5a9b602ccbcfdf7f016fc8da0596e", size = 124213 },
+]
+
+[[package]]
+name = "virtualenv"
+version = "20.26.6"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "distlib" },
+    { name = "filelock" },
+    { name = "importlib-metadata", marker = "python_full_version < '3.8'" },
+    { name = "platformdirs" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/3f/40/abc5a766da6b0b2457f819feab8e9203cbeae29327bd241359f866a3da9d/virtualenv-20.26.6.tar.gz", hash = "sha256:280aede09a2a5c317e409a00102e7077c6432c5a38f0ef938e643805a7ad2c48", size = 9372482 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/59/90/57b8ac0c8a231545adc7698c64c5a36fa7cd8e376c691b9bde877269f2eb/virtualenv-20.26.6-py3-none-any.whl", hash = "sha256:7345cc5b25405607a624d8418154577459c3e0277f5466dd79c49d5e492995f2", size = 5999862 },
+]
+
+[[package]]
+name = "werkzeug"
+version = "2.2.3"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "markupsafe" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/02/3c/baaebf3235c87d61d6593467056d5a8fba7c75ac838b8d100a5e64eba7a0/Werkzeug-2.2.3.tar.gz", hash = "sha256:2e1ccc9417d4da358b9de6f174e3ac094391ea1d4fbef2d667865d819dfd0afe", size = 845884 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/f6/f8/9da63c1617ae2a1dec2fbf6412f3a0cfe9d4ce029eccbda6e1e4258ca45f/Werkzeug-2.2.3-py3-none-any.whl", hash = "sha256:56433961bc1f12533306c624f3be5e744389ac61d722175d543e1751285da612", size = 233551 },
+]
+
+[[package]]
+name = "z3-solver"
+version = "4.15.3.0"
+source = { registry = "https://pypi.org/simple" }
+dependencies = [
+    { name = "importlib-resources", marker = "python_full_version < '3.9'" },
+]
+sdist = { url = "https://files.pythonhosted.org/packages/a3/60/9a924ee28cd1d12f2482834581d9024bf05110aa1098c056e847f05f7f76/z3_solver-4.15.3.0.tar.gz", hash = "sha256:78f69aebda5519bfd8af146a129f36cf4721a3c2667e80d9fe35cc9bb4d214a6", size = 4985945 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/63/45/dd8e9d7500faa05eafa589cc8f0f0b982ce575d51b455d62dab8e19dd571/z3_solver-4.15.3.0-py3-none-macosx_13_0_arm64.whl", hash = "sha256:65335aab295ded7c0ce27c85556067087a87052389ff160777d1a1d48ef0d74f", size = 36882388 },
+    { url = "https://files.pythonhosted.org/packages/54/9e/a11186061d9fead8be43bad7c75055585694124b2ccdd896ef249fe5824f/z3_solver-4.15.3.0-py3-none-macosx_13_0_x86_64.whl", hash = "sha256:3e62e93adff2def3537ff1ca67c3d58a6ca6d1944e0b5e774f88627b199d50e7", size = 39637842 },
+    { url = "https://files.pythonhosted.org/packages/b9/0b/f15168475e5493ea44fa3c5e642903f05d2b870db71ad05662ed87a06976/z3_solver-4.15.3.0-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a9afd9ceb290482097474d43f08415bcc1874f433189d1449f6c1508e9c68384", size = 29056003 },
+    { url = "https://files.pythonhosted.org/packages/a7/04/32a97b1f04175ec56213168ee3659709e876e3feacacb891ed3c26c1a82c/z3_solver-4.15.3.0-py3-none-manylinux_2_34_aarch64.whl", hash = "sha256:f61ef44552489077eedd7e6d9bed52ef1875decf86d66027742099a2703b1c77", size = 27074143 },
+    { url = "https://files.pythonhosted.org/packages/75/77/da54076a584557ea34f20800c68f725fe61f1dd987493fcb410b4a26f99f/z3_solver-4.15.3.0-py3-none-win32.whl", hash = "sha256:0c603f6bad7423d6411adda6af55030b725e3d30f54ea91b714abcedd73b848a", size = 13123666 },
+    { url = "https://files.pythonhosted.org/packages/c1/59/abc1bad8b25e9c576484ba65ca5ed225c8ed24d601ec242712f1c370b693/z3_solver-4.15.3.0-py3-none-win_amd64.whl", hash = "sha256:06abdf6c36f97c463aea827533504fd59476d015a65cf170a88bd6a53ba13ab5", size = 16203065 },
+]
+
+[[package]]
+name = "zipp"
+version = "3.15.0"
+source = { registry = "https://pypi.org/simple" }
+sdist = { url = "https://files.pythonhosted.org/packages/00/27/f0ac6b846684cecce1ee93d32450c45ab607f65c2e0255f0092032d91f07/zipp-3.15.0.tar.gz", hash = "sha256:112929ad649da941c23de50f356a2b5570c954b65150642bccdd66bf194d224b", size = 18454 }
+wheels = [
+    { url = "https://files.pythonhosted.org/packages/5b/fa/c9e82bbe1af6266adf08afb563905eb87cab83fde00a0a08963510621047/zipp-3.15.0-py3-none-any.whl", hash = "sha256:48904fc76a60e542af151aded95726c1a5c34ed43ab4134b597665c86d7ad556", size = 6758 },
+]

From 0644e64c79036ef1e260315b7d2feabba282bde7 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Fri, 3 Oct 2025 17:56:31 -0400
Subject: [PATCH 39/90] fix simple ndarray indexing (and also add mgrid)

---
 nki-examples/test_kernels.py | 78 ++++++++++++++++++++++--------------
 triton_viz/core/nki.py       | 31 +++++++-------
 2 files changed, 65 insertions(+), 44 deletions(-)

diff --git a/nki-examples/test_kernels.py b/nki-examples/test_kernels.py
index a3e05657..88bcc200 100644
--- a/nki-examples/test_kernels.py
+++ b/nki-examples/test_kernels.py
@@ -1,5 +1,4 @@
 import math
-import tp
 import neuronxcc.nki.language as nl
 import neuronxcc.nki as nki
 import numpy as np
@@ -81,6 +80,17 @@ def tmp0_kernel(a):  # works
     # for store(dst, value, mask), dst.shape, value.shape, mask.shape need to be same, return shape
     return out
 
+def mgrid_kernel(a):  # test nl.mgrid functionality
+    B, D = a.shape
+    out = nl.ndarray(a.shape, dtype=a.dtype, buffer=nl.shared_hbm)
+
+    coords = nl.mgrid[:B, :D]
+    mask = (coords[0] < B) & (coords[1] < D)
+    a_tmp = nl.load(a[coords[0], coords[1]], mask=mask)
+    nl.store(out, value=2*a_tmp, mask=mask) # TODO: masked_store
+
+    return out
+
 def xyz_kernel(a):  # works
     B, T, C, H, W = a.shape
     out = nl.ndarray(a.shape, dtype=a.dtype, buffer=nl.shared_hbm)
@@ -127,7 +137,7 @@ def tmp1_kernel(a):
 blocks_x = math.ceil(B / 128)
 blocks_y = math.ceil(D / 512)
 
-kernel = tmp1_kernel
+kernel = mgrid_kernel
 if kernel == add_kernel:
     kernel_grid = (blocks_x, blocks_y)
     kernel_args = (x.numpy(), y.numpy())
@@ -172,39 +182,49 @@ def tmp1_kernel(a):
     kernel_grid = (blocks_x, blocks_y)
     kernel_args = (x.numpy(),)
     z1 = x
+if kernel == mgrid_kernel:
+    B, D = 1024, 1024
+    x = torch.rand((B, D))
 
-TRITON_VIZ = False
+    # Expected result: simple copy (nl.mgrid is just tested with device_print)
+    z1 = 2 * x  # Should be identical to input
+
+    kernel_grid = (1, 1, 1)
+    kernel_args = (x.numpy(),)
+
+TRITON_VIZ = True
 
 if TRITON_VIZ:
     kernel = triton_viz.trace(clients=Tracer(), backend="nki")(kernel)
     kk = kernel[kernel_grid]
     z2 = kk(*kernel_args)
-    z2 = torch.from_numpy(z2)
-    print((z1 - z2).abs().max())
-
-    print(f"Number of launches: {len(launches)}")
-    if launches:
-        launch = launches[-1]
-        print(f"Number of records: {len(launch.records)}")
-        for i, record in enumerate(launch.records):
-            print(f"Record {i}: {type(record).__name__}")
-            if hasattr(record, "ptr"):
-                print(f"  ptr: {record.ptr}")
-            if hasattr(record, "offsets"):
-                print(f"  offsets shape: {record.offsets.shape}")
-            if hasattr(record, "masks"):
-                print(f"  masks shape: {record.masks.shape}")
-
-    # Try to launch visualization
-    try:
-        triton_viz.launch(share=False)
-    except Exception as e:
-        print(f"\nError during visualization: {e}")
-        import traceback
-
-        traceback.print_exc()
+    #print((z1 - z2).abs().max())
+
+    #print(f"Number of launches: {len(launches)}")
+    #if launches:
+    #    launch = launches[-1]
+    #    print(f"Number of records: {len(launch.records)}")
+    #    for i, record in enumerate(launch.records):
+    #        print(f"Record {i}: {type(record).__name__}")
+    #        if hasattr(record, "ptr"):
+    #            print(f"  ptr: {record.ptr}")
+    #        if hasattr(record, "offsets"):
+    #            print(f"  offsets shape: {record.offsets.shape}")
+    #        if hasattr(record, "masks"):
+    #            print(f"  masks shape: {record.masks.shape}")
+
+    ## Try to launch visualization
+    #try:
+    #    triton_viz.launch(share=False)
+    #except Exception as e:
+    #    print(f"\nError during visualization: {e}")
+    #    import traceback
+
+    #    traceback.print_exc()
 else:
     kernel = nki.jit(kernel)
     z2 = nki.simulate_kernel(kernel[kernel_grid], *kernel_args)
-    z2 = torch.from_numpy(z2)
-    print((z1 - z2).abs().max())
+
+z2 = torch.from_numpy(z2)
+print((z1 - z2).abs().max())
+assert torch.allclose(z1, z2)
diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index 0088001f..76898dbd 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -129,23 +129,24 @@ def __getitem__(self, keys):
         """Implement slicing operations for NDArray"""
         if self._value is None:
             raise AttributeError("NDArray has no value to slice")
+        if not isinstance(keys, tuple):
+            keys = (keys,)
 
         # Apply the slicing to the underlying numpy array
         new_keys = []
-        if isinstance(keys, tuple):
-            arr_dim = 0
-            for k in keys:
-                if isinstance(k, NDArray):
-                    dim_len = self._value.shape[arr_dim]
-                    new_keys.append(k._value.clip(0, dim_len - 1))
-                elif isinstance(k, NLSlice):
-                    new_keys.append(slice(k.start, k.stop, k.step))
-                elif k is None:
-                    new_keys.append(k)
-                    arr_dim -= 1 # add new dim -> revisit arr_dim for next key
-                else:
-                    new_keys.append(k)
-                arr_dim += 1
+        arr_dim = 0
+        for k in keys:
+            if isinstance(k, NDArray):
+                dim_len = self._value.shape[arr_dim]
+                new_keys.append(k._value.clip(0, dim_len - 1))
+            elif isinstance(k, NLSlice):
+                new_keys.append(slice(k.start, k.stop, k.step))
+            elif k is None:
+                new_keys.append(k)
+                arr_dim -= 1 # add new dim -> revisit arr_dim for next key
+            else:
+                new_keys.append(k)
+            arr_dim += 1
         
         sliced_value = self._value[tuple(new_keys)]
 
@@ -382,7 +383,7 @@ def patch():
     nl.affine_range = nki_builder.range
     nl.par_dim
     nl.zeros = nki_builder.zeros
-    nl.mgrid
+    nl.mgrid = NDArray(value=np.mgrid, buffer=nl.sbuf, name='mgrid')
     nl.matmul = nki_builder.matmul
     nl.copy = nki_builder.copy
 

From 2037047df22ae7950f47cb5335634058b9238b25 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Wed, 15 Oct 2025 21:06:26 -0400
Subject: [PATCH 40/90] masked load stuff (masked store still doesn't work
 right)

---
 nki-examples/test_kernels.py         |   6 +-
 tests/test_masked_load.py            | 615 +++++++++++++++++++++++++++
 triton_viz/core/__init__.py          |   6 +
 triton_viz/core/nki.py               |  43 ++
 triton_viz/core/nki_extract_slice.py | 177 ++++++++
 triton_viz/core/nki_masked_load.py   | 170 ++++++++
 6 files changed, 1014 insertions(+), 3 deletions(-)
 create mode 100644 tests/test_masked_load.py
 create mode 100644 triton_viz/core/nki_extract_slice.py
 create mode 100644 triton_viz/core/nki_masked_load.py

diff --git a/nki-examples/test_kernels.py b/nki-examples/test_kernels.py
index 88bcc200..48f439db 100644
--- a/nki-examples/test_kernels.py
+++ b/nki-examples/test_kernels.py
@@ -87,7 +87,7 @@ def mgrid_kernel(a):  # test nl.mgrid functionality
     coords = nl.mgrid[:B, :D]
     mask = (coords[0] < B) & (coords[1] < D)
     a_tmp = nl.load(a[coords[0], coords[1]], mask=mask)
-    nl.store(out, value=2*a_tmp, mask=mask) # TODO: masked_store
+    nl.store(out[coords[0], coords[1]], value=2*a_tmp, mask=mask)
 
     return out
 
@@ -126,7 +126,7 @@ def tmp1_kernel(a):
     a_tmp3 = nl.load(a[80:83, 80:85], mask=iy3)
     nl.device_print(f"a_tmp3:", a_tmp3)
 
-    nl.store(out[ix, iy], value=a_tmp, mask=mask)
+    #nl.store(out[ix, iy], value=a_tmp, mask=mask)
     return out
 
 
@@ -137,7 +137,7 @@ def tmp1_kernel(a):
 blocks_x = math.ceil(B / 128)
 blocks_y = math.ceil(D / 512)
 
-kernel = mgrid_kernel
+kernel = tmp1_kernel
 if kernel == add_kernel:
     kernel_grid = (blocks_x, blocks_y)
     kernel_args = (x.numpy(), y.numpy())
diff --git a/tests/test_masked_load.py b/tests/test_masked_load.py
new file mode 100644
index 00000000..40c6323c
--- /dev/null
+++ b/tests/test_masked_load.py
@@ -0,0 +1,615 @@
+import numpy as np
+from triton_viz.core.nki_masked_load import masked_load, masked_store
+
+def test_masked_load():
+    """Test the masked_load function with various scenarios."""
+
+    print("=== Testing masked_load function ===\n")
+
+    # Test 1: mask=None case (direct indexing)
+    print("Test 1: mask=None case")
+    arr = np.array([1, 2, 3, 4, 5])
+    result = masked_load(arr, (slice(1, 4),), mask=None)
+    expected = arr[1:4]
+    print(f"Input: {arr}")
+    print(f"Keys: slice(1, 4)")
+    print(f"Result: {result}")
+    print(f"Expected: {expected}")
+    print(f"Match: {np.array_equal(result, expected)}\n")
+
+    # Test 2: In-bounds indexing with mask
+    print("Test 2: In-bounds indexing with mask")
+    arr = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
+    mask = np.array([[True, True, False], [True, False, True], [False, True, True]])
+    result = masked_load(arr, (slice(0, 3), slice(0, 3)), mask=mask)
+    result[[0,1,2],[2,1,0]] = 0
+    expected = np.array([[1, 2, 0], [4, 0, 6], [0, 8, 9]])
+    print(f"Input array:\n{arr}")
+    print(f"Mask:\n{mask}")
+    print(f"Keys: (slice(0, 2), slice(0, 2))")
+    print(f"Result:\n{result}")
+    print(f"Expected:\n{expected}")
+    print(f"Match: {np.array_equal(result, expected)}\n")
+
+    # Test 3: Out-of-bounds indexing with mask
+    print("Test 3: Out-of-bounds indexing with mask")
+    arr = np.array([[1, 2], [3, 4]])
+    mask = np.array([
+        [True , False, False],
+        [False, True , False],
+        [False, False, False],
+        [False, False, False],
+        [False, False, False],
+    ])  # Shape doesn't match exactly, but demonstrates concept
+    # This would normally cause IndexError for arr[5], but with mask should work
+    result = masked_load(arr, (slice(0, 5), slice(0, 3)), mask=mask)
+    print(f"Input array:\n{arr}")
+    print(f"Keys: (slice(0, 5), slice(0, 3)) - exceeds array bounds")
+    print(f"Result shape: {result.shape}")
+    print(f"Result:\n{result}")
+    print("Success: Handled OOB indexing\n")
+
+    # Test 4: Integer indexing OOB
+    print("Test 4: Integer indexing out of bounds")
+    arr = np.array([10, 20, 30])
+    mask = np.array([False])  # Extra mask element
+    result = masked_load(arr, ([5],), mask=mask)  # Index 5 is OOB
+    print(f"Input: {arr}")
+    print(f"Keys: (5,) - OOB index")
+    print(f"Result: {result}")
+    print(f"Result type: {type(result)}")
+    print("Success: Handled OOB integer indexing\n")
+
+    # Test 5: Array indexing
+    print("Test 5: Array indexing")
+    arr = np.array([100, 200, 300, 400, 500])
+    indices = np.array([0, 2, 4])  # Valid indices
+    mask = np.array([True, True, True])
+    result = masked_load(arr, (indices,), mask=mask)
+    expected = arr[indices]
+    print(f"Input: {arr}")
+    print(f"Indices: {indices}")
+    print(f"Result: {result}")
+    print(f"Expected: {expected}")
+    print(f"Match: {np.array_equal(result, expected)}\n")
+
+    # Test 6: Mixed indexing types
+    print("Test 6: Mixed indexing types")
+    arr = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])
+    mask = np.array([True, False, True, False])
+    # Mix of slice and integer
+    result = masked_load(arr, (slice(0, 4), 2), mask=mask)  # Get column 2, allow extending rows
+    print(f"Input array:\n{arr}")
+    print(f"Keys: (slice(0, 4), 2) - slice extends beyond array rows")
+    print(f"Result: {result}")
+    print(f"Result shape: {result.shape}")
+    print("Success: Handled mixed indexing\n")
+
+    # Test 7: Mixed indexing types
+    print("Test 7: Mixed indexing types")
+    arr = np.array([
+        [
+            [1, 2, 3],
+            [4, 5, 6],
+            [7, 8, 9],
+        ],
+        [
+            [11, 12, 13],
+            [14, 15, 16],
+            [17, 18, 19],
+        ],
+        [
+            [21, 22, 23],
+            [24, 25, 26],
+            [27, 28, 29],
+        ],
+    ])
+    arr_slice = (
+        slice(None, None, None),
+        None,
+        np.arange(20)[:, None],
+        np.arange(3)[None, :],
+        None,
+        None
+    )
+    mask = np.mgrid[:3, :1, :20, :3, :1, :1]
+    mask = (mask[0] > 0) & (mask[2] < 3) & (mask[3] < 2)
+    result = masked_load(arr, arr_slice, mask=mask)
+    expected = np.array([
+        [
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+        ],
+        [
+            [11, 12,  6],
+            [14, 15,  6],
+            [17, 18,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+        ],
+        [
+            [21, 22,  6],
+            [24, 25,  6],
+            [27, 28,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            ],
+        ]
+    )[:, None, :, :, None, None]
+    assert np.allclose(result, expected)
+
+    # Test 8: incomplete slices
+    print("Test 8: incomplete slices")
+    arr = np.array([
+        [
+            [1, 2, 3],
+            [4, 5, 6],
+            [7, 8, 9],
+        ],
+        [
+            [11, 12, 13],
+            [14, 15, 16],
+            [17, 18, 19],
+        ],
+        [
+            [21, 22, 23],
+            [24, 25, 26],
+            [27, 28, 29],
+        ],
+    ])
+    arr_slice = (
+        slice(None, None, None),
+    )
+    mask = np.mgrid[:3, :3, :3]
+    mask = (0 < mask[0]) & (mask[0] < 2) & (mask[1] < 2) & (mask[2] != 1)
+    result = masked_load(arr, arr_slice, mask=mask)
+    expected = np.array([
+        [
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+        ],
+        [
+            [11,  6, 13],
+            [14,  6, 16],
+            [ 6,  6,  6],
+        ],
+        [
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+            [ 6,  6,  6],
+        ],
+    ])
+    assert np.allclose(result, expected)
+
+    print("Test 9: Make sure index error still happen when mask=True")
+    arr = np.array([
+        [
+            [1, 2, 3],
+            [4, 5, 6],
+            [7, 8, 9],
+        ],
+        [
+            [11, 12, 13],
+            [14, 15, 16],
+            [17, 18, 19],
+        ],
+        [
+            [21, 22, 23],
+            [24, 25, 26],
+            [27, 28, 29],
+        ],
+    ])
+    arr_slice = (
+        slice(None, None, None),
+        slice(0, 4)
+    )
+    mask = np.mgrid[:3, :4, :3]
+    mask = (mask[0] < 3) & (mask[1] < 3) & (mask[2] < 3)
+    mask[-1, -1, -1] = True # index error
+    try:
+        result = masked_load(arr, arr_slice, mask=mask)
+        raise RuntimeError("Should've raised an index error but did not")
+    except IndexError:
+        pass
+
+    print("Test 10: Make sure error if mask shape wrong")
+    arr = np.array([
+        [
+            [1, 2, 3],
+            [4, 5, 6],
+            [7, 8, 9],
+        ],
+        [
+            [11, 12, 13],
+            [14, 15, 16],
+            [17, 18, 19],
+        ],
+        [
+            [21, 22, 23],
+            [24, 25, 26],
+            [27, 28, 29],
+        ],
+    ])
+    arr_slice = (slice(None, None, None),)
+    mask = np.mgrid[:3, :4, :3]
+    mask = (mask[0] < 3) & (mask[1] < 3) & (mask[2] < 3)
+    try:
+        result = masked_load(arr, arr_slice, mask=mask)
+        raise RuntimeError("Should've raised an assertion error (arr[arr_slice].shape != mask.shape) but did not")
+    except AssertionError:
+        pass
+
+
+    print("Test 10: Make sure error if mask shape wrong")
+    arr = np.array([
+        [
+            [1, 2, 3],
+            [4, 5, 6],
+            [7, 8, 9],
+        ],
+        [
+            [11, 12, 13],
+            [14, 15, 16],
+            [17, 18, 19],
+        ],
+        [
+            [21, 22, 23],
+            [24, 25, 26],
+            [27, 28, 29],
+        ],
+    ])
+    arr_slice = (slice(None, None, None),)
+    mask = np.arange(27).reshape(3, 3, 3) % 2 == 0
+    result = masked_load(arr, arr_slice, mask=mask)
+    print(result)
+
+def test_masked_store():
+    """Test the masked_store function with various scenarios."""
+
+    print("=== Testing masked_store function ===\n")
+
+    # Test 1: mask=None case (direct indexing)
+    print("Test 1: mask=None case")
+    arr = np.array([1, 2, 3, 4, 5])
+    values = np.array([10, 20, 30])
+    arr_copy = arr.copy()
+    masked_store(arr_copy, (slice(1, 4),), values, mask=None)
+    expected = arr.copy()
+    expected[1:4] = values
+    print(f"Original: {arr}")
+    print(f"Values: {values}")
+    print(f"Result: {arr_copy}")
+    print(f"Expected: {expected}")
+    assert np.array_equal(arr_copy, expected)
+
+    # Test 2: In-bounds indexing with mask
+    print("Test 2: In-bounds indexing with mask")
+    arr = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
+    values = np.array([[100, 200, 0], [400, 0, 600], [0, 800, 900]])
+    mask = np.array([[True, True, False], [True, False, True], [False, True, True]])
+    arr_copy = arr.copy()
+    masked_store(arr_copy, (slice(0, 3), slice(0, 3)), values, mask=mask)
+    expected = np.array([[100, 200, 3], [400, 5, 600], [7, 800, 900]])
+    print(f"Original array:\n{arr}")
+    print(f"Values:\n{values}")
+    print(f"Mask:\n{mask}")
+    print(f"Result:\n{arr_copy}")
+    print(f"Expected:\n{expected}")
+    assert np.array_equal(arr_copy, expected)
+
+    # Test 3: Out-of-bounds indexing with mask (should succeed where mask=False)
+    print("Test 3: Out-of-bounds indexing with mask")
+    arr = np.array([[1, 2], [3, 4]])
+    values = np.array([
+        [100, 0, 0],
+        [0, 200, 0],
+        [0, 0, 0],
+        [0, 0, 0],
+        [0, 0, 0],
+    ])
+    mask = np.array([
+        [True , False, False],
+        [False, True , False],
+        [False, False, False],
+        [False, False, False],
+        [False, False, False],
+    ])
+    arr_copy = arr.copy()
+    masked_store(arr_copy, (slice(0, 5), slice(0, 3)), values, mask=mask)
+    expected = np.array([[100, 2], [3, 200]])
+    print(f"Original array:\n{arr}")
+    print(f"Keys: (slice(0, 5), slice(0, 3)) - exceeds array bounds")
+    print(f"Result:\n{arr_copy}")
+    print(f"Expected:\n{expected}")
+    assert np.array_equal(arr_copy, expected)
+
+    # Test 4: Integer indexing OOB
+    print("Test 4: Integer indexing out of bounds")
+    arr = np.array([10, 20, 30])
+    values = np.array([999])
+    mask = np.array([False])  # mask=False at OOB index
+    arr_copy = arr.copy()
+    masked_store(arr_copy, ([5],), values, mask=mask)  # Index 5 is OOB
+    print(f"Original: {arr}")
+    print(f"Keys: (5,) - OOB index")
+    print(f"Result: {arr_copy}")
+    print(f"Expected: {arr} (unchanged)")
+    assert np.array_equal(arr_copy, arr)
+
+    # Test 5: Array indexing
+    print("Test 5: Array indexing")
+    arr = np.array([100, 200, 300, 400, 500])
+    indices = np.array([0, 2, 4])
+    values = np.array([111, 333, 555])
+    mask = np.array([True, True, True])
+    arr_copy = arr.copy()
+    masked_store(arr_copy, (indices,), values, mask=mask)
+    expected = np.array([111, 200, 333, 400, 555])
+    print(f"Original: {arr}")
+    print(f"Indices: {indices}")
+    print(f"Values: {values}")
+    print(f"Result: {arr_copy}")
+    print(f"Expected: {expected}")
+    assert np.array_equal(arr_copy, expected)
+
+    # Test 6: Mixed indexing types
+    print("Test 6: Mixed indexing types")
+    arr = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])
+    values = np.array([100, 200, 300, 400])
+    mask = np.array([True, False, True, False])
+    arr_copy = arr.copy()
+    masked_store(arr_copy, (slice(0, 4), 2), values, mask=mask)  # Set column 2
+    expected = np.array([[1, 2, 100, 4], [5, 6, 7, 8], [9, 10, 300, 12]])
+    print(f"Original array:\n{arr}")
+    print(f"Values: {values}")
+    print(f"Keys: (slice(0, 4), 2) - slice extends beyond array rows")
+    print(f"Result:\n{arr_copy}")
+    print(f"Expected:\n{expected}")
+    assert np.array_equal(arr_copy, expected)
+
+    # Test 7: Complex 3D indexing (simplified)
+    print("Test 7: Complex 3D indexing")
+    arr = np.array([
+        [
+            [1, 2, 3],
+            [4, 5, 6],
+            [7, 8, 9],
+        ],
+        [
+            [11, 12, 13],
+            [14, 15, 16],
+            [17, 18, 19],
+        ],
+        [
+            [21, 22, 23],
+            [24, 25, 26],
+            [27, 28, 29],
+        ],
+    ])
+    ## Simpler test case that uses basic slicing with mask
+    #arr_slice = (slice(0, 3), slice(0, 3))  # Just use 2D slicing for simplicity
+    #mask = np.array([
+    #    [True, True, False],
+    #    [False, True, True],
+    #    [True, False, False]
+    #])
+    #values = np.array([
+    #    [100, 200, 0],
+    #    [0, 300, 400],
+    #    [500, 0, 0]
+    #])
+    #arr_copy = arr[:, :, 0]  # Take first slice along last dimension
+    #arr_copy = arr_copy.copy()
+    #masked_store(arr_copy, arr_slice, values, mask=mask)
+
+    #expected_modified = np.array([
+    #    [100, 200, 7],   # Original row 0: [1, 4, 7], mask: [True, True, False] -> [100, 200, 7]
+    #    [11, 300, 400],  # Original row 1: [11, 14, 17], mask: [False, True, True] -> [11, 300, 400]
+    #    [500, 24, 27]    # Original row 2: [21, 24, 27], mask: [True, False, False] -> [500, 24, 27]
+    #])
+
+    arr_slice = (
+        slice(None, None, None),
+        None,
+        np.arange(20)[:, None],
+        np.arange(3)[None, :],
+        None,
+        None
+    )
+    mask = np.mgrid[:3, :1, :20, :3, :1, :1]
+    mask = (mask[0] > 0) & (mask[2] < 3) & (mask[3] < 2)
+    result = masked_load(arr, arr_slice, mask=mask)
+    values = np.array([
+        [
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+        ],
+        [
+            [31, 32, 99],
+            [34, 35, 99],
+            [37, 38, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+        ],
+        [
+            [41, 42, 99],
+            [44, 45, 99],
+            [47, 48, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            [99, 99, 99],
+            ],
+        ]
+    )[:, None, :, :, None, None]
+    arr_copy = arr.copy()
+    masked_store(arr_copy, arr_slice, values, mask=mask)
+
+    expected_modified = np.array([
+        [
+            [1, 2, 3],
+            [4, 5, 6],
+            [7, 8, 9],
+        ],
+        [
+            [31, 32, 13],
+            [34, 35, 16],
+            [37, 38, 19],
+        ],
+        [
+            [41, 42, 23],
+            [44, 45, 26],
+            [47, 48, 29],
+        ],
+    ])
+
+    print(f"Input array shape: {arr_copy.shape}")
+    print(f"Keys: 2D slicing")
+    print(f"Values shape: {values.shape}")
+    print(f"Mask shape: {mask.shape}")
+    assert np.array_equal(arr_copy, expected_modified)
+
+    # Test 8: Error case - OOB with mask=True should raise IndexError
+    print("Test 8: OOB with mask=True should raise IndexError")
+    arr = np.array([[1, 2], [3, 4]])
+    values = np.array([[100, 200], [300, 400], [500, 600]])
+    mask = np.array([
+        [True, True],
+        [True, True],
+        [True, False],  # This position is OOB and mask=True
+    ])
+    arr_copy = arr.copy()
+    try:
+        masked_store(arr_copy, (slice(0, 3), slice(0, 2)), values, mask=mask)
+        raise RuntimeError("Should've raised an index error but did not")
+    except IndexError:
+        print("Success: Correctly raised IndexError for OOB with mask=True\n")
+
+    # Test 9: Error case - values and mask shape mismatch
+    print("Test 9: Values and mask shape mismatch should raise AssertionError")
+    arr = np.array([1, 2, 3])
+    values = np.array([10, 20, 30])  # Shape (3,)
+    mask = np.array([True, False])   # Shape (2,) - different from values
+    arr_copy = arr.copy()
+    try:
+        masked_store(arr_copy, (slice(0, 3),), values, mask=mask)
+        print(f"ERROR: Should have raised AssertionError but didn't")
+        print(f"values.shape: {values.shape}")
+        print(f"mask.shape: {mask.shape}")
+        raise RuntimeError("Should've raised an assertion error but did not")
+    except AssertionError:
+        print("Success: Correctly raised AssertionError for values/mask shape mismatch\n")
+
+    # Test 10: Values shape mismatch
+    print("Test 10: Values shape mismatch should raise AssertionError")
+    arr = np.array([1, 2, 3])
+    values = np.array([10, 20, 30, 40])  # Wrong shape
+    mask = np.array([True, False, True, False])
+    arr_copy = arr.copy()
+    try:
+        masked_store(arr_copy, (slice(0, 3),), values, mask=mask)
+        raise RuntimeError("Should've raised an assertion error but did not")
+    except AssertionError:
+        print("Success: Correctly raised AssertionError for values shape mismatch\n")
+
+    print("All masked_store tests completed!")
+
+if __name__ == "__main__":
+    test_masked_load()
+    print("\n" + "="*50 + "\n")
+    test_masked_store()
diff --git a/triton_viz/core/__init__.py b/triton_viz/core/__init__.py
index e6a5e16b..1f66dd83 100644
--- a/triton_viz/core/__init__.py
+++ b/triton_viz/core/__init__.py
@@ -26,6 +26,8 @@
     Rsqrt,
     CastImpl,
 )
+from .nki_masked_load import masked_load, masked_store
+from .nki_extract_slice import StoreCallTransformer, transform_code
 
 __all__ = [
     "trace",
@@ -55,4 +57,8 @@
     "Idiv",
     "Rsqrt",
     "CastImpl",
+    "masked_load",
+    "masked_store",
+    "StoreCallTransformer",
+    "transform_code",
 ]
diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index 76898dbd..ebc4e276 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -305,6 +305,45 @@ def store(self, dst: NDArray, value: NDArray, *, mask=None, **kwargs):
         dst._value[mask._value] = value._value.ravel()
         return dst
 
+    def _convert_keys_to_numpy(self, keys):
+        """Convert any NDArrays in keys to numpy arrays."""
+        if isinstance(keys, (tuple, list)):
+            return tuple(self._convert_keys_to_numpy(k) for k in keys)
+        elif isinstance(keys, NDArray):
+            return keys._value
+        else:
+            return keys
+
+    def masked_load(self, src: NDArray, keys, *, mask=None, **kwargs):
+        """Load array elements with masking for out-of-bounds errors."""
+        # Convert NDArray to numpy array
+        ndarray = src._value
+        mask_value = getattr(mask, "_value", mask) if mask is not None else None
+
+        # Convert any NDArrays in keys to numpy arrays
+        numpy_keys = self._convert_keys_to_numpy(keys)
+
+        # Call the actual masked_load function
+        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)
+
+    def masked_store(self, dst: NDArray, keys, value: NDArray, *, mask=None, **kwargs):
+        """Store array elements with masking for out-of-bounds errors."""
+        # Convert NDArrays to numpy arrays
+        ndarray = dst._value
+        value_array = value._value
+        mask_value = getattr(mask, "_value", mask) if mask is not None else None
+
+        # Convert any NDArrays in keys to numpy arrays
+        numpy_keys = self._convert_keys_to_numpy(keys)
+
+        # Call the actual masked_store function
+        masked_store(ndarray, numpy_keys, value_array, mask=mask_value)
+
+        return dst
+
     def unary_op(self, x: NDArray, np_func, op_name, **kwargs):
         return NDArray(value=np_func(x._value), name=f"{x.name}_{op_name}", **kwargs)
 
@@ -374,6 +413,10 @@ def patch():
     nl.arange = nki_builder.arange
     nl.load = nki_builder.load
     nl.store = nki_builder.store
+
+    # Also expose masked_load and masked_store functions
+    nl.masked_load = nki_builder.masked_load
+    nl.masked_store = nki_builder.masked_store
     # see https://awsdocs-neuron.readthedocs-hosted.com/en/latest/general/nki/api/nki.language.html
 
     # TODO: implement
diff --git a/triton_viz/core/nki_extract_slice.py b/triton_viz/core/nki_extract_slice.py
new file mode 100644
index 00000000..7b1c5267
--- /dev/null
+++ b/triton_viz/core/nki_extract_slice.py
@@ -0,0 +1,177 @@
+import ast
+
+class StoreCallTransformer(ast.NodeTransformer):
+    """
+    A targeted AST transformer to rewrite `nl.store(x[...], ...)` calls
+    into `masked_store(x, slice_obj, ...)` and `nl.load(x[...])` calls
+    into `masked_load(x, slice_obj)` with a preceding assignment.
+    """
+
+    # MODIFICATION 1: We now visit the 'Expr' statement wrapper.
+    def visit_Expr(self, node: ast.Expr) -> ast.AST | list[ast.AST]:
+        """
+        Intercept and transform expression statements.
+        """
+        # We only care about expressions that are function calls.
+        if not isinstance(node.value, ast.Call):
+            return self.generic_visit(node)
+
+        call_node = node.value
+
+        # --- The logic from the old visit_Call starts here ---
+
+        # 1. IDENTIFY: Is this the specific call we want to transform?
+        if not (
+            isinstance(call_node.func, ast.Attribute)
+            and isinstance(call_node.func.value, ast.Name)
+            and call_node.func.value.id == 'nl'
+            and call_node.func.attr in ('store', 'load')
+        ):
+            return self.generic_visit(node)
+
+        # 2. VALIDATE: Does the call have the expected structure?
+        if not call_node.args or not isinstance(call_node.args[0], ast.Subscript):
+            return self.generic_visit(node)
+
+        # 3. DECONSTRUCT
+        subscript_node = call_node.args[0]
+        sliced_object = subscript_node.value
+        slice_content = subscript_node.slice
+        remaining_args = call_node.args[1:]
+
+        # 4. REBUILD
+        # Convert to nl.masked_load or nl.masked_store
+        func_name = 'masked_' + call_node.func.attr
+
+        new_call = ast.Call(
+            func=ast.Attribute(
+                value=ast.Name(id='nl', ctx=ast.Load()),
+                attr=func_name,
+                ctx=ast.Load()
+            ),
+            args=[
+                sliced_object,
+                self._create_slice_value_node(slice_content),
+                *remaining_args
+            ],
+            keywords=call_node.keywords
+        )
+
+        # 5. REPLACE
+        return ast.Expr(value=new_call)
+
+    def visit_Assign(self, node: ast.Assign) -> ast.AST:
+        """
+        Handle assignment statements where the value is an nl.load or nl.store call.
+        """
+        # Check if the value being assigned is a function call
+        if not isinstance(node.value, ast.Call):
+            return self.generic_visit(node)
+
+        call_node = node.value
+
+        # 1. IDENTIFY: Is this the specific call we want to transform?
+        if not (
+            isinstance(call_node.func, ast.Attribute)
+            and isinstance(call_node.func.value, ast.Name)
+            and call_node.func.value.id == 'nl'
+            and call_node.func.attr in ('store', 'load')
+        ):
+            return self.generic_visit(node)
+
+        # 2. VALIDATE: Does the call have the expected structure?
+        if not call_node.args or not isinstance(call_node.args[0], ast.Subscript):
+            return self.generic_visit(node)
+
+        # 3. DECONSTRUCT
+        subscript_node = call_node.args[0]
+        sliced_object = subscript_node.value
+        slice_content = subscript_node.slice
+        remaining_args = call_node.args[1:]
+
+        # 4. REBUILD
+        # Convert to nl.masked_load or nl.masked_store
+        func_name = 'masked_' + call_node.func.attr
+
+        new_call = ast.Call(
+            func=ast.Attribute(
+                value=ast.Name(id='nl', ctx=ast.Load()),
+                attr=func_name,
+                ctx=ast.Load()
+            ),
+            args=[
+                sliced_object,
+                self._create_slice_value_node(slice_content),
+                *remaining_args
+            ],
+            keywords=call_node.keywords
+        )
+
+        # 5. REPLACE: Create a new assignment with the transformed call
+        return ast.Assign(
+            targets=node.targets,
+            value=new_call,
+            type_comment=getattr(node, 'type_comment', None)
+        )
+
+    def _create_slice_value_node(self, node: ast.AST) -> ast.AST:
+        """
+        (This helper function is unchanged and correct)
+        Recursively transforms a slice's AST content into a constructible object.
+        """
+        match node:
+            case ast.Slice(lower, upper, step):
+                return ast.Call(
+                    func=ast.Name(id='slice', ctx=ast.Load()),
+                    args=[
+                        lower or ast.Constant(value=None),
+                        upper or ast.Constant(value=None),
+                        step or ast.Constant(value=None)
+                    ],
+                    keywords=[]
+                )
+            case ast.Tuple(elts):
+                return ast.Tuple(
+                    elts=[self._create_slice_value_node(e) for e in elts],
+                    ctx=ast.Load()
+                )
+            case _:
+                return node
+
+def transform_code(source_code: str) -> str:
+    """
+    Applies the StoreCallTransformer to a string of Python code.
+    """
+    tree = ast.parse(source_code)
+    transformer = StoreCallTransformer()
+    new_tree = transformer.visit(tree)
+    ast.fix_missing_locations(new_tree)
+    return ast.unparse(new_tree)
+
+# --- Input Code ---
+source_code = """
+import numpy as np
+
+# This load call should also be transformed
+sbuf_value = nl.load(x[2:4, None, ..., nl.arange(128)[None, :]], mask=mask)
+
+# This is the call we want to transform
+nl.store(x[2:4, None, ..., nl.arange(128)[None, :]], sbuf_value, mask=mask)
+
+
+# This call should be ignored
+other_lib.store(y[10:])
+
+# This call should also be ignored
+nl.some_other_func(z[:5])
+"""
+
+if __name__ == '__main__':
+    # --- Transformation ---
+    transformed_code = transform_code(source_code)
+
+    # --- Output ---
+    print("--- Original Code ---")
+    print(source_code)
+    print("\n--- Transformed Code ---")
+    print(transformed_code)
diff --git a/triton_viz/core/nki_masked_load.py b/triton_viz/core/nki_masked_load.py
new file mode 100644
index 00000000..1094756b
--- /dev/null
+++ b/triton_viz/core/nki_masked_load.py
@@ -0,0 +1,170 @@
+import numpy as np
+
+get = lambda x, default: default if x is None else x
+
+def normalize_slice(ndarray: np.ndarray, keys: tuple) -> tuple[tuple, tuple]:
+    """ Separate singleton dims (None) and add bounds to slices like :2, 2:, or :"""
+    singleton_dims = []
+    dim_idx = 0
+    new_keys = []
+    for key_dim, k in enumerate(keys):
+        if k is None:
+            singleton_dims.append(key_dim)
+            continue
+
+        arr_dim = ndarray.shape[dim_idx]
+        if isinstance(k, slice):
+            # add bounds to [:, :N, N:]-type slices
+            start = get(k.start, 0)
+            stop = get(k.stop, arr_dim)
+            step = get(k.step, 1)
+            if start < 0:
+                start = max(0, arr_dim + start)
+            if stop < 0:
+                stop = max(0, arr_dim + stop)
+            k = slice(start, stop, step)
+
+        new_keys.append(k)
+        dim_idx += 1
+    return tuple(new_keys), tuple(singleton_dims)
+
+def _calculate_target_shape(keys: tuple, original_shape: tuple[int, ...]) -> tuple[int, ...]:
+    """Calculate the shape needed for the padded array to handle indexing."""
+    target_shape = list(original_shape)
+    for i, (arr_dim, key) in enumerate(zip(original_shape, keys)):
+        if isinstance(key, slice):
+            target_shape[i] = max(arr_dim, key.stop)
+        elif isinstance(key, (int, np.integer)):
+            target_shape[i] = max(arr_dim, int(key) + 1)
+        elif isinstance(key, (list, np.ndarray)):
+            target_shape[i] = max(arr_dim, np.max(key) + 1)
+    return tuple(target_shape)
+
+def _get_valid_indices(keys: tuple, original_shape: tuple[int, ...], result_shape: tuple[int, ...]) -> np.ndarray:
+    """Get a boolean mask indicating which result indices are within original array bounds."""
+    coords = np.mgrid[[slice(0, s) for s in result_shape]]
+    valid_mask = np.ones(result_shape, dtype=bool)
+
+    for key_idx, (arr_dim, key) in enumerate(zip(original_shape, keys)):
+        if isinstance(key, slice):
+            start = get(key.start, 0)
+            step = get(key.step, 1)
+            original_coords = coords[key_idx] * step + start
+            valid_mask &= (0 <= original_coords) & (original_coords < arr_dim)
+        elif isinstance(key, (list, np.ndarray)):
+            valid_indices = np.array(key)
+            valid_mask &= (0 <= valid_indices) & (valid_indices < arr_dim)
+        elif isinstance(key, (int, np.integer)):
+            valid_mask &= (0 <= int(key) < arr_dim)
+
+    return valid_mask
+
+def masked_load(ndarray: np.ndarray, keys: tuple, mask: np.ndarray = None) -> np.ndarray:
+    """
+    Load array elements with masking for out-of-bounds errors.
+
+    Args:
+        ndarray: Input numpy array
+        keys: Indexing keys (tuple of slices, integers, arrays, etc.)
+        mask: Boolean mask array. If None, returns direct indexing.
+              Where mask is False at OOB indices, garbage values are kept.
+              Mask shape should match the shape of the result, not the input array.
+
+    Returns:
+        Indexed array with masked error handling
+    """
+    # Handle mask=None case
+    if mask is None:
+        return ndarray[keys]
+
+    # Convert keys to tuple if it's not already
+    if not isinstance(keys, tuple):
+        keys = (keys,)
+
+    # normalize (remove Nones in slices)
+    keys, singleton_dims = normalize_slice(ndarray, keys)
+    target_shape = _calculate_target_shape(keys, ndarray.shape)
+    padded_array = np.empty(target_shape, dtype=ndarray.dtype)
+
+    # Calculate the region where we can safely copy the original array
+    copy_slices = []
+    for i in range(len(ndarray.shape)):
+        copy_slices.append(slice(0, ndarray.shape[i]))
+    padded_array[tuple(copy_slices)] = ndarray[tuple(copy_slices)]
+
+    result = padded_array[keys]
+    assert np.expand_dims(result, singleton_dims).shape == mask.shape
+
+    mask = np.squeeze(mask, singleton_dims)
+
+    # Determine which indices are actually within original array bounds
+    in_bounds_mask = _get_valid_indices(keys, ndarray.shape, result.shape)
+
+    # Check if there are any OOB indices where mask=True
+    oob_with_true_mask = (~in_bounds_mask) & mask
+    oob_coords = np.where(oob_with_true_mask)
+    if len(oob_coords[0]) > 0:
+        # Get the first OOB coordinate
+        oob_idx = tuple(coord[0] for coord in oob_coords)
+        raise IndexError(f"index {oob_idx} is out of bounds for array of size {ndarray.shape}")
+
+    valid_mask = mask & in_bounds_mask
+    result[~valid_mask] = 6
+
+    return np.expand_dims(result, singleton_dims)
+
+def masked_store(ndarray: np.ndarray, keys: tuple, value: np.ndarray, mask: np.ndarray = None) -> None:
+    """
+    Store array elements with masking for out-of-bounds errors.
+
+    Args:
+        ndarray: Input numpy array to store values into
+        keys: Indexing keys (tuple of slices, integers, arrays, etc.)
+        value: Values to store
+        mask: Boolean mask array. If None, performs direct indexing.
+              Where mask is False at OOB indices, values are not stored.
+              Mask shape should match the shape of the result, not the input array.
+
+    Returns:
+        None (modifies ndarray in-place)
+    """
+    # Handle mask=None case
+    if mask is None:
+        import tp; tp.log(f'{ndarray.shape=}, {keys=}')
+        ndarray[keys] = value
+        return
+
+    assert value.shape == mask.shape
+
+    # Convert keys to tuple if it's not already
+    if not isinstance(keys, tuple):
+        keys = (keys,)
+
+    # normalize (remove Nones in slices)
+    keys, singleton_dims = normalize_slice(ndarray, keys)
+    target_shape = _calculate_target_shape(keys, ndarray.shape)
+    offset_array = np.arange(np.prod(target_shape)).reshape(target_shape)
+
+    offsets = np.expand_dims(offset_array[keys], singleton_dims)
+    assert offsets.shape == mask.shape
+
+    mask = np.squeeze(mask, singleton_dims)
+
+    # Check if there are any OOB indices where mask=True
+    in_bounds_mask = _get_valid_indices(keys, ndarray.shape, offsets.shape)
+    oob_with_true_mask = (~in_bounds_mask) & mask
+    oob_coords = np.where(oob_with_true_mask)
+    if len(oob_coords[0]) > 0:
+        # Get the first OOB coordinate
+        oob_idx = tuple(coord[0] for coord in oob_coords)
+        raise IndexError(f"index {oob_idx} is out of bounds for array of size {ndarray.shape}")
+
+    print('!Ndarray')
+    print(ndarray)
+    print('!values')
+    print(value)
+    print('!Offsets')
+    print(offsets)
+    print('!mask')
+    print(mask)
+    ndarray.ravel()[offsets[mask]] = value.ravel()[offsets[mask]]

From 8f619e8288f19d47d61fc7d5e4c18ce3100853ac Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Wed, 15 Oct 2025 21:08:16 -0400
Subject: [PATCH 41/90] attempt to connect nki interpreter to triton-viz (still
 failing)

---
 triton_viz/clients/tracer/tracer.py |   8 +-
 triton_viz/core/nki.py              | 101 +++++++++++++++++-
 triton_viz/core/patch.py            | 160 +++++++++++++++-------------
 3 files changed, 184 insertions(+), 85 deletions(-)

diff --git a/triton_viz/clients/tracer/tracer.py b/triton_viz/clients/tracer/tracer.py
index dd6e3292..995e1988 100644
--- a/triton_viz/clients/tracer/tracer.py
+++ b/triton_viz/clients/tracer/tracer.py
@@ -56,7 +56,7 @@ def grid_idx_callback(self, grid_idx: tuple[int]):
     def grid_callback(self, grid: tuple[int]):
         self.tensors = sorted(self.tensors, key=lambda x: x.data_ptr())
 
-    def register_op_callback(self, op_type: type[Op]) -> OpCallbacks:
+    def register_op_callback(self, op_type: type[Op], *ignore_args, **ignore_kwargs) -> OpCallbacks:
         def pre_load_callback(
             ptr, mask, *ignore_args, **ignore_kwargs
         ):
@@ -68,7 +68,7 @@ def pre_load_callback(
                 Load(tensor.data_ptr(), ptr.data - tensor.data_ptr(), mask.data)
             )
 
-        def pre_store_callback(ptr, value, mask, cache_modifier, eviction_policy):
+        def pre_store_callback(ptr, value, mask, *ignore_args, **ignore_kwargs):
             if not self.sample:
                 return
             first_ptr = np.reshape(ptr.data, (-1))[0]
@@ -77,14 +77,14 @@ def pre_store_callback(ptr, value, mask, cache_modifier, eviction_policy):
                 Store(tensor.data_ptr(), ptr.data - tensor.data_ptr(), mask.data)
             )
 
-        def post_reduce_sum_callback(ret, input, axis=None, keep_dims=False):
+        def post_reduce_sum_callback(ret, input, axis=None, keep_dims=False, *ignore_args, **ignore_kwargs):
             if not self.sample:
                 return
             input_shape = input.handle.data.shape
             output_shape = ret.handle.data.shape
             self.records.append(ReduceSum(input_shape, axis, keep_dims, output_shape))
 
-        def post_dot_callback(ret, input, other, *args):
+        def post_dot_callback(ret, input, other, *ignore_args, **ignore_kwargs):
             if not self.sample:
                 return
             input_shape = input.data.shape
diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index ebc4e276..cb837dad 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -2,6 +2,9 @@
 
 import neuronxcc.nki.language as nl
 import inspect
+import ast
+from .nki_extract_slice import StoreCallTransformer, transform_code
+from .nki_masked_load import masked_load, masked_store
 
 
 # Q1: slicing semantic is weird
@@ -120,7 +123,54 @@ def value(self, new_value):
         self._value = new_value
 
     def data_ptr(self):
-        return self.ctypes.data
+        return self._value.ctypes.data
+
+    def stride(self):
+        return self._value.strides
+    
+    def element_size(self):
+        return self.dtype.itemsize
+
+    def cpu(self): # THTODO: rm?
+        return self
+
+    def get_offsets(self):
+        """
+        Generate offset arrays for each dimension based on shape and stride.
+
+        Args:
+            strides: Tuple of strides for each dimension (a, b, ..., z)
+
+        Returns:
+            Tuple of offset arrays: (arange(A)[:, None, ..., None]*a, arange(B)[None, :, ..., None]*b, ...)
+        """
+        strides = self._value.strides
+        if self._value is None:
+            raise AttributeError("NDArray has no value - cannot compute offsets")
+
+        shape = self.shape
+        if len(shape) != len(strides):
+            raise ValueError(f"Shape has {len(shape)} dimensions but strides has {len(strides)} dimensions")
+
+        #offsets = []
+        offsets = 0
+        ndim = len(shape)
+
+        for i, (dim_size, stride) in enumerate(zip(shape, strides)):
+            # Create arange for this dimension
+            arange_vals = np.arange(dim_size)
+
+            # Create broadcast shape - put arange in position i, others as 1
+            broadcast_shape = [1] * ndim
+            broadcast_shape[i] = dim_size
+
+            # Reshape and multiply by stride
+            offset_array = (arange_vals * stride).reshape(broadcast_shape)
+            #offsets.append(NDArray(value=offset_array, name=f"{self.name}_offset_dim{i}"))
+            offsets += NDArray(value=offset_array, name=f"{self.name}_offset_dim{i}")
+
+        #return tuple(offsets)
+        return offsets
 
     def __repr__(self):
         return f"NDArray(shape={self.shape}, dtype={self.dtype}, name={self.name})"
@@ -451,6 +501,8 @@ def patch():
     nl.pow = nki_builder.pow
     nl.reciprocal = nki_builder.reciprocal
 
+    nl.device_print = print
+
 
 def unpatch():
     # reload the original functions
@@ -477,9 +529,43 @@ def run(self, *args, **kwargs):
         nki_builder.fn = self.fn
 
         kwargs.pop("warmup", None)  # Remove warmup from kwargs if it exists
-        kwargs.pop("client_manager", None)  # Remove client_manager from kwargs if it exists
-
-        patch()
+        client_manager = kwargs.pop("client_manager", None)  # Remove client_manager from kwargs if it exists
+
+        # Call grid_callback once before grid execution (similar to Triton)
+        if client_manager is not None:
+            client_manager.grid_callback(grid_dims)
+
+        patch() # NKI interpreter patching
+
+        #with client_manager.patch():
+        #v
+        #kwargs.update({"client_manager": client_manager})
+        #ret = self.interpreter_fn.run(*args, **kwargs)
+        #self.finalize()
+        #return ret
+
+        # Apply AST transformer to convert nl.load/nl.store calls to nl.masked_load/nl.masked_store
+        import types
+        if hasattr(self.fn, '__code__'):
+            # Get the source code of the function
+            source_code = inspect.getsource(self.fn)
+            # Transform the source code using the AST transformer
+            transformed_code = transform_code(source_code)
+            # Create a new function from the transformed code
+            #exec_globals = self.fn.__globals__.copy()
+            #exec(transformed_code, exec_globals)
+            ## Replace the original function with the transformed one
+            #self.fn = exec_globals[self.fn.__name__]
+            exec_globals = self.fn.__globals__.copy()
+            import random, string, os
+            rand_str = ''.join(random.choices(string.ascii_letters + string.digits, k=16))
+            os.makedirs('/tmp/triton-viz', exist_ok=True)
+            filename = f'/tmp/triton-viz/{rand_str}.py'
+            with open(filename, 'w') as f:
+                f.write(transformed_code)
+            code_obj = compile(transformed_code, filename=filename, mode="exec")
+            exec(code_obj, exec_globals)
+            self.fn = exec_globals[self.fn.__name__]
 
         # convert args to NDArray if they are not already
         args = [arg if isinstance(arg, NDArray) else NDArray(value=arg) for arg in args]
@@ -487,6 +573,13 @@ def run(self, *args, **kwargs):
             for y in range(grid_dims[1]):
                 for z in range(grid_dims[2]):
                     nki_builder.set_grid_idx(x, y, z)
+
+                    # Call grid_idx_callback for each grid iteration (similar to Triton)
+                    if client_manager is not None:
+                        client_manager.grid_idx_callback((x, y, z))
+
                     result = self.fn(*args, **kwargs)
+        #^
+
         unpatch()
         return result.value
diff --git a/triton_viz/core/patch.py b/triton_viz/core/patch.py
index c99a5563..7d1078cd 100644
--- a/triton_viz/core/patch.py
+++ b/triton_viz/core/patch.py
@@ -41,79 +41,85 @@
 from triton.runtime import JITFunction
 from triton_viz.core.nki import nki_builder
 
-op_list = [
-    ProgramId,
-    RawStore,
-    Store,
-    RawLoad,
-    Load,
-    UnaryOp,
-    BinaryOp,
-    TernaryOp,
-    Dot,
-    MakeRange,
-    AddPtr,
-    Splat,
-    ExpandDims,
-    Broadcast,
-    ReduceMax,
-    ReduceMin,
-    ReduceSum,
-    MakeBlockPointer,
-    TensorPointerLoad,
-    TensorPointerStore,
-    Idiv,
-    Rsqrt,
-    CastImpl,
-
-    #ProgramId,
-    #Store,
-    #Load,
-    #Dot,
-    #UnaryOp,
-    #MakeRange
-]
-original_ops = {
-    ProgramId: interpreter_builder.create_get_program_id,
-    RawStore: interpreter_builder.create_store,
-    Store: interpreter_builder.create_masked_store,
-    RawLoad: interpreter_builder.create_load,
-    Load: interpreter_builder.create_masked_load,
-    Dot: interpreter_builder.create_dot,
-    UnaryOp: interpreter_builder.unary_op,
-    BinaryOp: interpreter_builder.binary_op,
-    TernaryOp: interpreter_builder.ternary_op,
-    MakeRange: interpreter_builder.create_make_range,
-    AddPtr: interpreter_builder.create_addptr,
-    ExpandDims: interpreter_builder.create_expand_dims,
-    Broadcast: interpreter_builder.create_broadcast,
-    Splat: interpreter_builder.create_splat,
-    MakeBlockPointer: interpreter_builder.create_make_block_ptr,
-    TensorPointerLoad: interpreter_builder.create_tensor_pointer_load,
-    TensorPointerStore: interpreter_builder.create_tensor_pointer_store,
-    Idiv: interpreter_builder.create_idiv,
-    Rsqrt: interpreter_builder.create_rsqrt,
-    CastImpl: interpreter_builder.cast_impl,
+BUILDER = interpreter_builder
+if BUILDER == interpreter_builder:
+    op_list = [
+        ProgramId,
+        RawStore,
+        Store,
+        RawLoad,
+        Load,
+        UnaryOp,
+        BinaryOp,
+        TernaryOp,
+        Dot,
+        MakeRange,
+        AddPtr,
+        Splat,
+        ExpandDims,
+        Broadcast,
+        ReduceMax,
+        ReduceMin,
+        ReduceSum,
+        MakeBlockPointer,
+        TensorPointerLoad,
+        TensorPointerStore,
+        Idiv,
+        Rsqrt,
+        CastImpl,
+    ]
+    original_ops = {
+        ProgramId: interpreter_builder.create_get_program_id,
+        RawStore: interpreter_builder.create_store,
+        Store: interpreter_builder.create_masked_store,
+        RawLoad: interpreter_builder.create_load,
+        Load: interpreter_builder.create_masked_load,
+        Dot: interpreter_builder.create_dot,
+        UnaryOp: interpreter_builder.unary_op,
+        BinaryOp: interpreter_builder.binary_op,
+        TernaryOp: interpreter_builder.ternary_op,
+        MakeRange: interpreter_builder.create_make_range,
+        AddPtr: interpreter_builder.create_addptr,
+        ExpandDims: interpreter_builder.create_expand_dims,
+        Broadcast: interpreter_builder.create_broadcast,
+        Splat: interpreter_builder.create_splat,
+        MakeBlockPointer: interpreter_builder.create_make_block_ptr,
+        TensorPointerLoad: interpreter_builder.create_tensor_pointer_load,
+        TensorPointerStore: interpreter_builder.create_tensor_pointer_store,
+        Idiv: interpreter_builder.create_idiv,
+        Rsqrt: interpreter_builder.create_rsqrt,
+        CastImpl: interpreter_builder.cast_impl,
+    }
+elif BUILDER == nki_builder:
+    op_list = [
+        ProgramId,
+        Store,
+        Load,
+        Dot,
+        UnaryOp,
+        MakeRange
+    ]
+    original_ops = {
+        ProgramId: nki_builder.program_id,
+        Store: nki_builder.store,
+        Load: nki_builder.load,
+        Dot: nki_builder.matmul,
+        UnaryOp: nki_builder.unary_op,
+        #BinaryOp: nki_builder.binary_op,
+        #TernaryOp: nki_builder.ternary_op,
+        MakeRange: nki_builder.arange,
+        #AddPtr: nki_builder.create_addptr,
+        #ExpandDims: nki_builder.create_expand_dims,
+        #Broadcast: nki_builder.create_broadcast,
+        #Splat: nki_builder.create_splat,
+        #MakeBlockPointer: nki_builder.create_make_block_ptr,
+        #TensorPointerLoad: nki_builder.create_tensor_pointer_load,
+        #TensorPointerStore: nki_builder.create_tensor_pointer_store,
+        #Idiv: nki_builder.create_idiv,
+        #Rsqrt: nki_builder.create_rsqrt,
+        #CastImpl: nki_builder.cast_impl,
+    }
 
-    #ProgramId: nki_builder.program_id,
-    #Store: nki_builder.store,
-    #Load: nki_builder.load,
-    #Dot: nki_builder.matmul,
-    #UnaryOp: nki_builder.unary_op,
-    ##BinaryOp: nki_builder.binary_op,
-    ##TernaryOp: nki_builder.ternary_op,
-    #MakeRange: nki_builder.arange,
-    ##AddPtr: nki_builder.create_addptr,
-    ##ExpandDims: nki_builder.create_expand_dims,
-    ##Broadcast: nki_builder.create_broadcast,
-    ##Splat: nki_builder.create_splat,
-    ##MakeBlockPointer: nki_builder.create_make_block_ptr,
-    ##TensorPointerLoad: nki_builder.create_tensor_pointer_load,
-    ##TensorPointerStore: nki_builder.create_tensor_pointer_store,
-    ##Idiv: nki_builder.create_idiv,
-    ##Rsqrt: nki_builder.create_rsqrt,
-    ##CastImpl: nki_builder.cast_impl,
-}
 reduce_map: dict[type[Op], Callable] = {
     ReduceMax: tl.max,
     ReduceMin: tl.min,
@@ -168,9 +174,9 @@ def patch_op(op_type: type[Op], callbacks: OpCallbacks):
     if op_type in original_ops:
         # create a new function that calls the before_callback, the original op and the after_callback
         op_name = original_ops[op_type].__name__
-        current_op = getattr(interpreter_builder, op_name)
+        current_op = getattr(BUILDER, op_name)
         patched_op = PatchOp(current_op, op_type, callbacks)
-        setattr(interpreter_builder, op_name, patched_op)
+        setattr(BUILDER, op_name, patched_op)
     elif op_type in reduce_map:
         op_name = reduce_map[op_type].__name__
         current_op = getattr(tl, op_name)
@@ -188,7 +194,7 @@ def unpatch_op(op_type: type[Op]):
     """
     if op_type in original_ops:
         op_name = original_ops[op_type].__name__
-        setattr(interpreter_builder, op_name, original_ops[op_type])
+        setattr(BUILDER, op_name, original_ops[op_type])
 
 
 def _patch_lang(fn):
@@ -226,7 +232,7 @@ def run_grid_loops(grid):
                     leave=False,
                     disable=not (cfg.report_grid_execution_progress and grid[2] > 1),
                 ):
-                    interpreter_builder.set_grid_idx(x, y, z)
+                    BUILDER.set_grid_idx(x, y, z)
                     client_manager.grid_idx_callback((x, y, z))
                     self.fn(**call_args)
                     # if symbolic execution, only do one iteration
@@ -263,7 +269,7 @@ def run_grid_loops(grid):
     grid = self.grid(call_args) if callable(self.grid) else self.grid
     assert len(grid) <= 3
     grid = grid + (1,) * (3 - len(grid))
-    interpreter_builder.set_grid_dim(*grid)
+    BUILDER.set_grid_dim(*grid)
     client_manager.grid_callback(grid)
     run_grid_loops(grid)
     # Copy arguments back to propagate side-effects

From 5adb5e035e6c4c797193d95a339a316a642937bf Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Wed, 15 Oct 2025 21:08:58 -0400
Subject: [PATCH 42/90] hardcode offsets and callbacks for NKI matmul kernel

---
 nki-examples/matmul.py | 102 +++++++++++++++++++++++++++--------------
 1 file changed, 67 insertions(+), 35 deletions(-)

diff --git a/nki-examples/matmul.py b/nki-examples/matmul.py
index e0eab83d..4e3c458b 100644
--- a/nki-examples/matmul.py
+++ b/nki-examples/matmul.py
@@ -11,6 +11,8 @@
 import numpy as np
 import math
 
+tensors = [] # THTODO: rm
+records = [triton_viz.core.data.Grid((0,0,0))] # THTODO: rm
 def matmul_kernel(lhs, rhs):
     """NKI matmul_kernel to compute a matrix multiplication operation in a tiled manner
 
@@ -34,6 +36,9 @@ def matmul_kernel(lhs, rhs):
     TILE_K = 2
     TILE_N = 4
 
+    tensors.append(lhs)
+    tensors.append(rhs)
+    tensors.append(result)
     # Use affine_range to loop over tiles
     m_range = nl.arange(M)[:, None]
     n_range = nl.arange(N)[None, :]
@@ -50,62 +55,89 @@ def matmul_kernel(lhs, rhs):
                 # Load tiles from lhs and rhs
                 lhs_p = nl.arange(TILE_M)[:, None] + m * TILE_M
                 lhs_f = nl.arange(TILE_K)[None, :] + k * TILE_K
-                lhs_tile = nl.load(lhs[lhs_p, lhs_f], mask=(lhs_p < M) & (lhs_f < K))
+                lhs_mask = (lhs_p < M) & (lhs_f < K)
+                lhs_tile = nl.load(lhs[lhs_p, lhs_f], mask=lhs_mask)
+                records.append(triton_viz.core.data.Load(
+                    lhs.data_ptr(),
+                    (4 * (lhs_p * K + lhs_f))._value,
+                    lhs_mask._value
+                )) # THTODO: rm
+
                 rhs_p = nl.arange(TILE_K)[:, None] + k * TILE_K
                 rhs_f = nl.arange(TILE_N)[None, :] + n * TILE_N
-                rhs_tile = nl.load(rhs[rhs_p, rhs_f], mask=(rhs_p < K) & (rhs_f < N))
+                rhs_mask = (rhs_p < K) & (rhs_f < N) 
+                rhs_tile = nl.load(rhs[rhs_p, rhs_f], mask=rhs_mask)
+                records.append(triton_viz.core.data.Load(
+                    rhs.data_ptr(),
+                    (4 * (rhs_p * N + rhs_f))._value,
+                    rhs_mask._value
+                )) # THTODO: rm
 
                 # Accumulate partial-sums into PSUM
-                res_psum += nl.matmul(lhs_tile[...], rhs_tile[...], transpose_x=False)
+                x = nl.matmul(lhs_tile[...], rhs_tile[...], transpose_x=False)
+                res_psum += x
+                records.append(triton_viz.core.data.Dot(
+                    lhs_tile.shape,
+                    rhs_tile.shape,
+                    res_psum.shape,
+                    lhs_tile._value,
+                    rhs_tile._value,
+                    x._value
+                )) # THTODO: rm
 
             # Copy the result from PSUM back to SBUF, and cast to expected output data-type
-            out_mask = (m * TILE_M <= m_range) & (m_range < (m + 1) * TILE_M)
-            out_mask &= (n * TILE_N <= n_range) & (n_range < (n + 1) * TILE_N)
             res_sb = nl.copy(res_psum, dtype=result.dtype)
 
-            #nl.store(result[m * TILE_M:(m + 1) * TILE_M, n * TILE_N:(n + 1) * TILE_N],
-            #         value=res_sb) # TODO: uncomment this
-            nl.store(result, value=res_sb, mask=out_mask) # TODO: remove this - incorrect NKI syntax but currently needed for interpreter
+            out_p = nl.arange(TILE_M)[:, None] + m * TILE_M
+            out_f = nl.arange(TILE_N)[None, :] + n * TILE_N
+            out_mask = (out_p < M) & (out_f < N)
+            nl.store(result[m * TILE_M:(m + 1) * TILE_M, n * TILE_N:(n + 1) * TILE_N],
+                     value=res_sb)
+            records.append(triton_viz.core.data.Store(
+                result.data_ptr(),
+                (4 * (out_p * N + out_f))._value,
+                out_mask._value
+            )) # THTODO: rm
 
     return result
 
 TRITON_VIZ = True
 kernel_grid = (1, 1, 1)
-lhs_small = np.arange(64).astype(np.float32).reshape(8, 8)
-rhs_small = np.arange(128).astype(np.float32).reshape(8, 16)
-#lhs_small = np.arange(16).astype(np.float32).reshape(4, 4)
-#rhs_small = np.arange(32).astype(np.float32).reshape(4, 8)
+lhs_small = np.arange(16).astype(np.float32).reshape(4, 4)
+rhs_small = np.arange(32).astype(np.float32).reshape(4, 8)
 kernel_args = (lhs_small, rhs_small)
 
 if TRITON_VIZ:
-    tp.log('Executing matmul_kernel with NKI interpreter...')
+    print('Executing matmul_kernel with NKI interpreter...')
     matmul_kernel = triton_viz.trace(clients=Tracer(), backend="nki")(matmul_kernel)
     kk = matmul_kernel[kernel_grid]
     z2 = kk(*kernel_args)
 
-    #print(f"Number of launches: {len(launches)}")
-    #if launches:
-    #    launch = launches[-1]
-    #    print(f"Number of records: {len(launch.records)}")
-    #    for i, record in enumerate(launch.records):
-    #        print(f"Record {i}: {type(record).__name__}")
-    #        if hasattr(record, "ptr"):
-    #            print(f"  ptr: {record.ptr}")
-    #        if hasattr(record, "offsets"):
-    #            print(f"  offsets shape: {record.offsets.shape}")
-    #        if hasattr(record, "masks"):
-    #            print(f"  masks shape: {record.masks.shape}")
-
-    ## Try to launch visualization
-    #try:
-    #    triton_viz.launch(share=False)
-    #except Exception as e:
-    #    print(f"\nError during visualization: {e}")
-    #    import traceback
-
-    #    traceback.print_exc()
+    print(f"Number of launches: {len(launches)}")
+    if launches:
+        launch = launches[-1]
+        launch.records = records # THTODO: rm
+        launch.tensors = tensors # THTODO: rm
+        print(f"Number of records: {len(launch.records)}")
+        for i, record in enumerate(launch.records):
+            print(f"Record {i}: {type(record).__name__}")
+            if hasattr(record, "ptr"):
+                print(f"  ptr: {record.ptr}")
+            if hasattr(record, "offsets"):
+                print(f"  offsets shape: {record.offsets.shape}")
+            if hasattr(record, "masks"):
+                print(f"  masks shape: {record.masks.shape}")
+
+    # Try to launch visualization
+    try:
+        triton_viz.launch(share=False)
+    except Exception as e:
+        print(f"\nError during visualization: {e}")
+        import traceback
+
+        traceback.print_exc()
 else:
-    tp.log('Executing NKI JIT-ed matmul_kernel...')
+    print('Executing NKI JIT-ed matmul_kernel...')
     matmul_kernel = nki.jit(matmul_kernel)
     z2 = nki.simulate_kernel(matmul_kernel[kernel_grid], *kernel_args)
 

From f091fd4d4d5060b35d21de5c5cb58f5fffc670e6 Mon Sep 17 00:00:00 2001
From: Hao Wu <hwu27@gmu.edu>
Date: Thu, 16 Oct 2025 18:51:56 -0400
Subject: [PATCH 43/90] Remove uv.lock

---
 uv.lock | 795 --------------------------------------------------------
 1 file changed, 795 deletions(-)
 delete mode 100644 uv.lock

diff --git a/uv.lock b/uv.lock
deleted file mode 100644
index b2a5d4cb..00000000
--- a/uv.lock
+++ /dev/null
@@ -1,795 +0,0 @@
-version = 1
-requires-python = ">=3.7"
-
-[[package]]
-name = "anytree"
-version = "2.8.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "six" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/d8/45/de59861abc8cb66e9e95c02b214be4d52900aa92ce34241a957dcf1d569d/anytree-2.8.0.tar.gz", hash = "sha256:3f0f93f355a91bc3e6245319bf4c1d50e3416cc7a35cc1133c1ff38306bbccab", size = 189484 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/a8/65/be23d8c3ecd68d40541d49812cd94ed0f3ee37eb88669ca15df0e43daed1/anytree-2.8.0-py2.py3-none-any.whl", hash = "sha256:14c55ac77492b11532395049a03b773d14c7e30b22aa012e337b1e983de31521", size = 41717 },
-]
-
-[[package]]
-name = "cairocffi"
-version = "1.6.1"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "cffi" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/6e/37/982a8650ab8c0b3bd0021e3f7041e9707ce35f6cc43b7def2c43ff109dd9/cairocffi-1.6.1.tar.gz", hash = "sha256:78e6bbe47357640c453d0be929fa49cd05cce2e1286f3d2a1ca9cbda7efdb8b7", size = 87436 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/17/be/a5d2c16317c6a890502725970589ae7f06cfc66b2e6916ba0a86973403c8/cairocffi-1.6.1-py3-none-any.whl", hash = "sha256:aa78ee52b9069d7475eeac457389b6275aa92111895d78fbaa2202a52dac112e", size = 75096 },
-]
-
-[[package]]
-name = "certifi"
-version = "2025.8.3"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/dc/67/960ebe6bf230a96cda2e0abcf73af550ec4f090005363542f0765df162e0/certifi-2025.8.3.tar.gz", hash = "sha256:e564105f78ded564e3ae7c923924435e1daa7463faeab5bb932bc53ffae63407", size = 162386 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/e5/48/1549795ba7742c948d2ad169c1c8cdbae65bc450d6cd753d124b17c8cd32/certifi-2025.8.3-py3-none-any.whl", hash = "sha256:f6c12493cfb1b06ba2ff328595af9350c65d6644968e5d3a2ffd78699af217a5", size = 161216 },
-]
-
-[[package]]
-name = "cffi"
-version = "1.15.1"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "pycparser" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/2b/a8/050ab4f0c3d4c1b8aaa805f70e26e84d0e27004907c5b8ecc1d31815f92a/cffi-1.15.1.tar.gz", hash = "sha256:d400bfb9a37b1351253cb402671cea7e89bdecc294e8016a707f6d1d8ac934f9", size = 508501 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/e8/ff/c4b7a358526f231efa46a375c959506c87622fb4a2c5726e827c55e6adf2/cffi-1.15.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:39d39875251ca8f612b6f33e6b1195af86d1b3e60086068be9cc053aa4376e21", size = 179233 },
-    { url = "https://files.pythonhosted.org/packages/ea/be/c4ad40ad441ac847b67c7a37284ae3c58f39f3e638c6b0f85fb662233825/cffi-1.15.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:285d29981935eb726a4399badae8f0ffdff4f5050eaa6d0cfc3f64b857b77185", size = 174280 },
-    { url = "https://files.pythonhosted.org/packages/ed/a3/c5f01988ddb70a187c3e6112152e01696188c9f8a4fa4c68aa330adbb179/cffi-1.15.1-cp310-cp310-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:3eb6971dcff08619f8d91607cfc726518b6fa2a9eba42856be181c6d0d9515fd", size = 421712 },
-    { url = "https://files.pythonhosted.org/packages/ef/41/19da352d341963d29a33bdb28433ba94c05672fb16155f794fad3fd907b0/cffi-1.15.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:21157295583fe8943475029ed5abdcf71eb3911894724e360acff1d61c1d54bc", size = 449886 },
-    { url = "https://files.pythonhosted.org/packages/af/da/9441d56d7dd19d07dcc40a2a5031a1f51c82a27cee3705edf53dadcac398/cffi-1.15.1-cp310-cp310-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:5635bd9cb9731e6d4a1132a498dd34f764034a8ce60cef4f5319c0541159392f", size = 450520 },
-    { url = "https://files.pythonhosted.org/packages/aa/02/ab15b3aa572759df752491d5fa0f74128cd14e002e8e3257c1ab1587810b/cffi-1.15.1-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:2012c72d854c2d03e45d06ae57f40d78e5770d252f195b93f581acf3ba44496e", size = 446015 },
-    { url = "https://files.pythonhosted.org/packages/88/89/c34caf63029fb7628ec2ebd5c88ae0c9bd17db98c812e4065a4d020ca41f/cffi-1.15.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:dd86c085fae2efd48ac91dd7ccffcfc0571387fe1193d33b6394db7ef31fe2a4", size = 441830 },
-    { url = "https://files.pythonhosted.org/packages/32/bd/d0809593f7976828f06a492716fbcbbfb62798bbf60ea1f65200b8d49901/cffi-1.15.1-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:fa6693661a4c91757f4412306191b6dc88c1703f780c8234035eac011922bc01", size = 434743 },
-    { url = "https://files.pythonhosted.org/packages/0e/65/0d7b5dad821ced4dcd43f96a362905a68ce71e6b5f5cfd2fada867840582/cffi-1.15.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:59c0b02d0a6c384d453fece7566d1c7e6b7bae4fc5874ef2ef46d56776d61c9e", size = 464113 },
-    { url = "https://files.pythonhosted.org/packages/9f/52/1e2b43cfdd7d9a39f48bc89fcaee8d8685b1295e205a4f1044909ac14d89/cffi-1.15.1-cp310-cp310-win32.whl", hash = "sha256:cba9d6b9a7d64d4bd46167096fc9d2f835e25d7e4c121fb2ddfc6528fb0413b2", size = 170412 },
-    { url = "https://files.pythonhosted.org/packages/0e/e2/a23af3d81838c577571da4ff01b799b0c2bbde24bd924d97e228febae810/cffi-1.15.1-cp310-cp310-win_amd64.whl", hash = "sha256:ce4bcc037df4fc5e3d184794f27bdaab018943698f4ca31630bc7f84a7b69c6d", size = 179060 },
-    { url = "https://files.pythonhosted.org/packages/23/8b/2e8c2469eaf89f7273ac685164949a7e644cdfe5daf1c036564208c3d26b/cffi-1.15.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:3d08afd128ddaa624a48cf2b859afef385b720bb4b43df214f85616922e6a5ac", size = 179198 },
-    { url = "https://files.pythonhosted.org/packages/f9/96/fc9e118c47b7adc45a0676f413b4a47554e5f3b6c99b8607ec9726466ef1/cffi-1.15.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:3799aecf2e17cf585d977b780ce79ff0dc9b78d799fc694221ce814c2c19db83", size = 174221 },
-    { url = "https://files.pythonhosted.org/packages/10/72/617ee266192223a38b67149c830bd9376b69cf3551e1477abc72ff23ef8e/cffi-1.15.1-cp311-cp311-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:a591fe9e525846e4d154205572a029f653ada1a78b93697f3b5a8f1f2bc055b9", size = 441694 },
-    { url = "https://files.pythonhosted.org/packages/91/bc/b7723c2fe7a22eee71d7edf2102cd43423d5f95ff3932ebaa2f82c7ec8d0/cffi-1.15.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3548db281cd7d2561c9ad9984681c95f7b0e38881201e157833a2342c30d5e8c", size = 470613 },
-    { url = "https://files.pythonhosted.org/packages/5d/4e/4e0bb5579b01fdbfd4388bd1eb9394a989e1336203a4b7f700d887b233c1/cffi-1.15.1-cp311-cp311-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:91fc98adde3d7881af9b59ed0294046f3806221863722ba7d8d120c575314325", size = 472199 },
-    { url = "https://files.pythonhosted.org/packages/37/5a/c37631a86be838bdd84cc0259130942bf7e6e32f70f4cab95f479847fb91/cffi-1.15.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:94411f22c3985acaec6f83c6df553f2dbe17b698cc7f8ae751ff2237d96b9e3c", size = 462588 },
-    { url = "https://files.pythonhosted.org/packages/71/d7/0fe0d91b0bbf610fb7254bb164fa8931596e660d62e90fb6289b7ee27b09/cffi-1.15.1-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:03425bdae262c76aad70202debd780501fabeaca237cdfddc008987c0e0f59ef", size = 450543 },
-    { url = "https://files.pythonhosted.org/packages/d3/56/3e94aa719ae96eeda8b68b3ec6e347e0a23168c6841dc276ccdcdadc9f32/cffi-1.15.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:cc4d65aeeaa04136a12677d3dd0b1c0c94dc43abac5860ab33cceb42b801c1e8", size = 474253 },
-    { url = "https://files.pythonhosted.org/packages/87/ee/ddc23981fc0f5e7b5356e98884226bcb899f95ebaefc3e8e8b8742dd7e22/cffi-1.15.1-cp311-cp311-win32.whl", hash = "sha256:a0f100c8912c114ff53e1202d0078b425bee3649ae34d7b070e9697f93c5d52d", size = 170313 },
-    { url = "https://files.pythonhosted.org/packages/43/a0/cc7370ef72b6ee586369bacd3961089ab3d94ae712febf07a244f1448ffd/cffi-1.15.1-cp311-cp311-win_amd64.whl", hash = "sha256:04ed324bda3cda42b9b695d51bb7d54b680b9719cfab04227cdd1e04e5de3104", size = 179001 },
-    { url = "https://files.pythonhosted.org/packages/b5/7d/df6c088ef30e78a78b0c9cca6b904d5abb698afb5bc8f5191d529d83d667/cffi-1.15.1-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:198caafb44239b60e252492445da556afafc7d1e3ab7a1fb3f0584ef6d742375", size = 178906 },
-    { url = "https://files.pythonhosted.org/packages/c2/0b/3b09a755ddb977c167e6d209a7536f6ade43bb0654bad42e08df1406b8e4/cffi-1.15.1-cp37-cp37m-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:5ef34d190326c3b1f822a5b7a45f6c4535e2f47ed06fec77d3d799c450b2651e", size = 405992 },
-    { url = "https://files.pythonhosted.org/packages/5b/1a/e1ee5bed11d8b6540c05a8e3c32448832d775364d4461dd6497374533401/cffi-1.15.1-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8102eaf27e1e448db915d08afa8b41d6c7ca7a04b7d73af6514df10a3e74bd82", size = 435560 },
-    { url = "https://files.pythonhosted.org/packages/d3/e1/e55ca2e0dd446caa2cc8f73c2b98879c04a1f4064ac529e1836683ca58b8/cffi-1.15.1-cp37-cp37m-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:5df2768244d19ab7f60546d0c7c63ce1581f7af8b5de3eb3004b9b6fc8a9f84b", size = 435478 },
-    { url = "https://files.pythonhosted.org/packages/2e/7a/68c35c151e5b7a12650ecc12fdfb85211aa1da43e9924598451c4a0a3839/cffi-1.15.1-cp37-cp37m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:a8c4917bd7ad33e8eb21e9a5bbba979b49d9a97acb3a803092cbc1133e20343c", size = 430395 },
-    { url = "https://files.pythonhosted.org/packages/93/d0/2e2b27ea2f69b0ec9e481647822f8f77f5fc23faca2dd00d1ff009940eb7/cffi-1.15.1-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:0e2642fe3142e4cc4af0799748233ad6da94c62a8bec3a6648bf8ee68b1c7426", size = 427911 },
-    { url = "https://files.pythonhosted.org/packages/50/34/4cc590ad600869502c9838b4824982c122179089ed6791a8b1c95f0ff55e/cffi-1.15.1-cp37-cp37m-win32.whl", hash = "sha256:e229a521186c75c8ad9490854fd8bbdd9a0c9aa3a524326b55be83b54d4e0ad9", size = 169721 },
-    { url = "https://files.pythonhosted.org/packages/32/2a/63cb8c07d151de92ff9d897b2eb27ba6a0e78dda8e4c5f70d7b8c16cd6a2/cffi-1.15.1-cp37-cp37m-win_amd64.whl", hash = "sha256:a0b71b1b8fbf2b96e41c4d990244165e2c9be83d54962a9a1d118fd8657d2045", size = 179301 },
-    { url = "https://files.pythonhosted.org/packages/87/4b/64e8bd9d15d6b22b6cb11997094fbe61edf453ea0a97c8675cb7d1c3f06f/cffi-1.15.1-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:320dab6e7cb2eacdf0e658569d2575c4dad258c0fcc794f46215e1e39f90f2c3", size = 178940 },
-    { url = "https://files.pythonhosted.org/packages/22/c6/df826563f55f7e9dd9a1d3617866282afa969fe0d57decffa1911f416ed8/cffi-1.15.1-cp38-cp38-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:1e74c6b51a9ed6589199c787bf5f9875612ca4a8a0785fb2d4a84429badaf22a", size = 421947 },
-    { url = "https://files.pythonhosted.org/packages/c1/25/16a082701378170559bb1d0e9ef2d293cece8dc62913d79351beb34c5ddf/cffi-1.15.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a5c84c68147988265e60416b57fc83425a78058853509c1b0629c180094904a5", size = 449906 },
-    { url = "https://files.pythonhosted.org/packages/df/02/aef53d4aa43154b829e9707c8c60bab413cd21819c4a36b0d7aaa83e2a61/cffi-1.15.1-cp38-cp38-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:3b926aa83d1edb5aa5b427b4053dc420ec295a08e40911296b9eb1b6170f6cca", size = 451028 },
-    { url = "https://files.pythonhosted.org/packages/79/4b/33494eb0adbcd884656c48f6db0c98ad8a5c678fb8fb5ed41ab546b04d8c/cffi-1.15.1-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:87c450779d0914f2861b8526e035c5e6da0a3199d8f1add1a665e1cbc6fc6d02", size = 446520 },
-    { url = "https://files.pythonhosted.org/packages/b7/8b/06f30caa03b5b3ac006de4f93478dbd0239e2a16566d81a106c322dc4f79/cffi-1.15.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:4f2c9f67e9821cad2e5f480bc8d83b8742896f1242dba247911072d4fa94c192", size = 442655 },
-    { url = "https://files.pythonhosted.org/packages/47/97/137f0e3d2304df2060abb872a5830af809d7559a5a4b6a295afb02728e65/cffi-1.15.1-cp38-cp38-win32.whl", hash = "sha256:8b7ee99e510d7b66cdb6c593f21c043c248537a32e0bedf02e01e9553a172314", size = 170236 },
-    { url = "https://files.pythonhosted.org/packages/c9/e3/0a52838832408cfbbf3a59cb19bcd17e64eb33795c9710ca7d29ae10b5b7/cffi-1.15.1-cp38-cp38-win_amd64.whl", hash = "sha256:00a9ed42e88df81ffae7a8ab6d9356b371399b91dbdf0c3cb1e84c03a13aceb5", size = 178835 },
-    { url = "https://files.pythonhosted.org/packages/18/8f/5ff70c7458d61fa8a9752e5ee9c9984c601b0060aae0c619316a1e1f1ee5/cffi-1.15.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:54a2db7b78338edd780e7ef7f9f6c442500fb0d41a5a4ea24fff1c929d5af585", size = 179249 },
-    { url = "https://files.pythonhosted.org/packages/3a/75/a162315adeaf47e94a3b7f886a8e31d77b9e525a387eef2d6f0efc96a7c8/cffi-1.15.1-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:fcd131dd944808b5bdb38e6f5b53013c5aa4f334c5cad0c72742f6eba4b73db0", size = 174297 },
-    { url = "https://files.pythonhosted.org/packages/85/1f/a3c533f8d377da5ca7edb4f580cc3edc1edbebc45fac8bb3ae60f1176629/cffi-1.15.1-cp39-cp39-manylinux_2_12_i686.manylinux2010_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:7473e861101c9e72452f9bf8acb984947aa1661a7704553a9f6e4baa5ba64415", size = 420641 },
-    { url = "https://files.pythonhosted.org/packages/77/b7/d3618d612be01e184033eab90006f8ca5b5edafd17bf247439ea4e167d8a/cffi-1.15.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6c9a799e985904922a4d207a94eae35c78ebae90e128f0c4e521ce339396be9d", size = 448814 },
-    { url = "https://files.pythonhosted.org/packages/a9/ba/e082df21ebaa9cb29f2c4e1d7e49a29b90fcd667d43632c6674a16d65382/cffi-1.15.1-cp39-cp39-manylinux_2_17_ppc64le.manylinux2014_ppc64le.whl", hash = "sha256:3bcde07039e586f91b45c88f8583ea7cf7a0770df3a1649627bf598332cb6984", size = 449647 },
-    { url = "https://files.pythonhosted.org/packages/af/cb/53b7bba75a18372d57113ba934b27d0734206c283c1dfcc172347fbd9f76/cffi-1.15.1-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:33ab79603146aace82c2427da5ca6e58f2b3f2fb5da893ceac0c42218a40be35", size = 445191 },
-    { url = "https://files.pythonhosted.org/packages/2d/86/3ca57cddfa0419f6a95d1c8478f8f622ba597e3581fd501bbb915b20eb75/cffi-1.15.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5d598b938678ebf3c67377cdd45e09d431369c3b1a5b331058c338e201f12b27", size = 441236 },
-    { url = "https://files.pythonhosted.org/packages/ad/26/7b3a73ab7d82a64664c7c4ea470e4ec4a3c73bb4f02575c543a41e272de5/cffi-1.15.1-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:db0fbb9c62743ce59a9ff687eb5f4afbe77e5e8403d6697f7446e5f609976f76", size = 433865 },
-    { url = "https://files.pythonhosted.org/packages/da/ff/ab939e2c7b3f40d851c0f7192c876f1910f3442080c9c846532993ec3cef/cffi-1.15.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:98d85c6a2bef81588d9227dde12db8a7f47f639f4a17c9ae08e773aa9c697bf3", size = 463090 },
-    { url = "https://files.pythonhosted.org/packages/c6/3d/dd085bb831b22ce4d0b7ba8550e6d78960f02f770bbd1314fea3580727f8/cffi-1.15.1-cp39-cp39-win32.whl", hash = "sha256:40f4774f5a9d4f5e344f31a32b5096977b5d48560c5592e2f3d2c4374bd543ee", size = 170382 },
-    { url = "https://files.pythonhosted.org/packages/a8/16/06b84a7063a4c0a2b081030fdd976022086da9c14e80a9ed4ba0183a98a9/cffi-1.15.1-cp39-cp39-win_amd64.whl", hash = "sha256:70df4e3b545a17496c9b3f41f5115e69a4f2e77e94e1d2a8e1070bc0c38c8a3c", size = 179079 },
-]
-
-[[package]]
-name = "cfgv"
-version = "3.3.1"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/c4/bf/d0d622b660d414a47dc7f0d303791a627663f554345b21250e39e7acb48b/cfgv-3.3.1.tar.gz", hash = "sha256:f5a830efb9ce7a445376bb66ec94c638a9787422f96264c98edc6bdeed8ab736", size = 7864 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/6d/82/0a0ebd35bae9981dea55c06f8e6aaf44a49171ad798795c72c6f64cba4c2/cfgv-3.3.1-py2.py3-none-any.whl", hash = "sha256:c6a0883f3917a037485059700b9e75da2464e6c27051014ad85ba6aaa5884426", size = 7312 },
-]
-
-[[package]]
-name = "charset-normalizer"
-version = "3.4.3"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/83/2d/5fd176ceb9b2fc619e63405525573493ca23441330fcdaee6bef9460e924/charset_normalizer-3.4.3.tar.gz", hash = "sha256:6fce4b8500244f6fcb71465d4a4930d132ba9ab8e71a7859e6a5d59851068d14", size = 122371 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/d6/98/f3b8013223728a99b908c9344da3aa04ee6e3fa235f19409033eda92fb78/charset_normalizer-3.4.3-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:fb7f67a1bfa6e40b438170ebdc8158b78dc465a5a67b6dde178a46987b244a72", size = 207695 },
-    { url = "https://files.pythonhosted.org/packages/21/40/5188be1e3118c82dcb7c2a5ba101b783822cfb413a0268ed3be0468532de/charset_normalizer-3.4.3-cp310-cp310-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:cc9370a2da1ac13f0153780040f465839e6cccb4a1e44810124b4e22483c93fe", size = 147153 },
-    { url = "https://files.pythonhosted.org/packages/37/60/5d0d74bc1e1380f0b72c327948d9c2aca14b46a9efd87604e724260f384c/charset_normalizer-3.4.3-cp310-cp310-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:07a0eae9e2787b586e129fdcbe1af6997f8d0e5abaa0bc98c0e20e124d67e601", size = 160428 },
-    { url = "https://files.pythonhosted.org/packages/85/9a/d891f63722d9158688de58d050c59dc3da560ea7f04f4c53e769de5140f5/charset_normalizer-3.4.3-cp310-cp310-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:74d77e25adda8581ffc1c720f1c81ca082921329452eba58b16233ab1842141c", size = 157627 },
-    { url = "https://files.pythonhosted.org/packages/65/1a/7425c952944a6521a9cfa7e675343f83fd82085b8af2b1373a2409c683dc/charset_normalizer-3.4.3-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:d0e909868420b7049dafd3a31d45125b31143eec59235311fc4c57ea26a4acd2", size = 152388 },
-    { url = "https://files.pythonhosted.org/packages/f0/c9/a2c9c2a355a8594ce2446085e2ec97fd44d323c684ff32042e2a6b718e1d/charset_normalizer-3.4.3-cp310-cp310-musllinux_1_2_aarch64.whl", hash = "sha256:c6f162aabe9a91a309510d74eeb6507fab5fff92337a15acbe77753d88d9dcf0", size = 150077 },
-    { url = "https://files.pythonhosted.org/packages/3b/38/20a1f44e4851aa1c9105d6e7110c9d020e093dfa5836d712a5f074a12bf7/charset_normalizer-3.4.3-cp310-cp310-musllinux_1_2_ppc64le.whl", hash = "sha256:4ca4c094de7771a98d7fbd67d9e5dbf1eb73efa4f744a730437d8a3a5cf994f0", size = 161631 },
-    { url = "https://files.pythonhosted.org/packages/a4/fa/384d2c0f57edad03d7bec3ebefb462090d8905b4ff5a2d2525f3bb711fac/charset_normalizer-3.4.3-cp310-cp310-musllinux_1_2_s390x.whl", hash = "sha256:02425242e96bcf29a49711b0ca9f37e451da7c70562bc10e8ed992a5a7a25cc0", size = 159210 },
-    { url = "https://files.pythonhosted.org/packages/33/9e/eca49d35867ca2db336b6ca27617deed4653b97ebf45dfc21311ce473c37/charset_normalizer-3.4.3-cp310-cp310-musllinux_1_2_x86_64.whl", hash = "sha256:78deba4d8f9590fe4dae384aeff04082510a709957e968753ff3c48399f6f92a", size = 153739 },
-    { url = "https://files.pythonhosted.org/packages/2a/91/26c3036e62dfe8de8061182d33be5025e2424002125c9500faff74a6735e/charset_normalizer-3.4.3-cp310-cp310-win32.whl", hash = "sha256:d79c198e27580c8e958906f803e63cddb77653731be08851c7df0b1a14a8fc0f", size = 99825 },
-    { url = "https://files.pythonhosted.org/packages/e2/c6/f05db471f81af1fa01839d44ae2a8bfeec8d2a8b4590f16c4e7393afd323/charset_normalizer-3.4.3-cp310-cp310-win_amd64.whl", hash = "sha256:c6e490913a46fa054e03699c70019ab869e990270597018cef1d8562132c2669", size = 107452 },
-    { url = "https://files.pythonhosted.org/packages/7f/b5/991245018615474a60965a7c9cd2b4efbaabd16d582a5547c47ee1c7730b/charset_normalizer-3.4.3-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:b256ee2e749283ef3ddcff51a675ff43798d92d746d1a6e4631bf8c707d22d0b", size = 204483 },
-    { url = "https://files.pythonhosted.org/packages/c7/2a/ae245c41c06299ec18262825c1569c5d3298fc920e4ddf56ab011b417efd/charset_normalizer-3.4.3-cp311-cp311-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:13faeacfe61784e2559e690fc53fa4c5ae97c6fcedb8eb6fb8d0a15b475d2c64", size = 145520 },
-    { url = "https://files.pythonhosted.org/packages/3a/a4/b3b6c76e7a635748c4421d2b92c7b8f90a432f98bda5082049af37ffc8e3/charset_normalizer-3.4.3-cp311-cp311-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:00237675befef519d9af72169d8604a067d92755e84fe76492fef5441db05b91", size = 158876 },
-    { url = "https://files.pythonhosted.org/packages/e2/e6/63bb0e10f90a8243c5def74b5b105b3bbbfb3e7bb753915fe333fb0c11ea/charset_normalizer-3.4.3-cp311-cp311-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:585f3b2a80fbd26b048a0be90c5aae8f06605d3c92615911c3a2b03a8a3b796f", size = 156083 },
-    { url = "https://files.pythonhosted.org/packages/87/df/b7737ff046c974b183ea9aa111b74185ac8c3a326c6262d413bd5a1b8c69/charset_normalizer-3.4.3-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:0e78314bdc32fa80696f72fa16dc61168fda4d6a0c014e0380f9d02f0e5d8a07", size = 150295 },
-    { url = "https://files.pythonhosted.org/packages/61/f1/190d9977e0084d3f1dc169acd060d479bbbc71b90bf3e7bf7b9927dec3eb/charset_normalizer-3.4.3-cp311-cp311-musllinux_1_2_aarch64.whl", hash = "sha256:96b2b3d1a83ad55310de8c7b4a2d04d9277d5591f40761274856635acc5fcb30", size = 148379 },
-    { url = "https://files.pythonhosted.org/packages/4c/92/27dbe365d34c68cfe0ca76f1edd70e8705d82b378cb54ebbaeabc2e3029d/charset_normalizer-3.4.3-cp311-cp311-musllinux_1_2_ppc64le.whl", hash = "sha256:939578d9d8fd4299220161fdd76e86c6a251987476f5243e8864a7844476ba14", size = 160018 },
-    { url = "https://files.pythonhosted.org/packages/99/04/baae2a1ea1893a01635d475b9261c889a18fd48393634b6270827869fa34/charset_normalizer-3.4.3-cp311-cp311-musllinux_1_2_s390x.whl", hash = "sha256:fd10de089bcdcd1be95a2f73dbe6254798ec1bda9f450d5828c96f93e2536b9c", size = 157430 },
-    { url = "https://files.pythonhosted.org/packages/2f/36/77da9c6a328c54d17b960c89eccacfab8271fdaaa228305330915b88afa9/charset_normalizer-3.4.3-cp311-cp311-musllinux_1_2_x86_64.whl", hash = "sha256:1e8ac75d72fa3775e0b7cb7e4629cec13b7514d928d15ef8ea06bca03ef01cae", size = 151600 },
-    { url = "https://files.pythonhosted.org/packages/64/d4/9eb4ff2c167edbbf08cdd28e19078bf195762e9bd63371689cab5ecd3d0d/charset_normalizer-3.4.3-cp311-cp311-win32.whl", hash = "sha256:6cf8fd4c04756b6b60146d98cd8a77d0cdae0e1ca20329da2ac85eed779b6849", size = 99616 },
-    { url = "https://files.pythonhosted.org/packages/f4/9c/996a4a028222e7761a96634d1820de8a744ff4327a00ada9c8942033089b/charset_normalizer-3.4.3-cp311-cp311-win_amd64.whl", hash = "sha256:31a9a6f775f9bcd865d88ee350f0ffb0e25936a7f930ca98995c05abf1faf21c", size = 107108 },
-    { url = "https://files.pythonhosted.org/packages/e9/5e/14c94999e418d9b87682734589404a25854d5f5d0408df68bc15b6ff54bb/charset_normalizer-3.4.3-cp312-cp312-macosx_10_13_universal2.whl", hash = "sha256:e28e334d3ff134e88989d90ba04b47d84382a828c061d0d1027b1b12a62b39b1", size = 205655 },
-    { url = "https://files.pythonhosted.org/packages/7d/a8/c6ec5d389672521f644505a257f50544c074cf5fc292d5390331cd6fc9c3/charset_normalizer-3.4.3-cp312-cp312-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:0cacf8f7297b0c4fcb74227692ca46b4a5852f8f4f24b3c766dd94a1075c4884", size = 146223 },
-    { url = "https://files.pythonhosted.org/packages/fc/eb/a2ffb08547f4e1e5415fb69eb7db25932c52a52bed371429648db4d84fb1/charset_normalizer-3.4.3-cp312-cp312-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:c6fd51128a41297f5409deab284fecbe5305ebd7e5a1f959bee1c054622b7018", size = 159366 },
-    { url = "https://files.pythonhosted.org/packages/82/10/0fd19f20c624b278dddaf83b8464dcddc2456cb4b02bb902a6da126b87a1/charset_normalizer-3.4.3-cp312-cp312-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:3cfb2aad70f2c6debfbcb717f23b7eb55febc0bb23dcffc0f076009da10c6392", size = 157104 },
-    { url = "https://files.pythonhosted.org/packages/16/ab/0233c3231af734f5dfcf0844aa9582d5a1466c985bbed6cedab85af9bfe3/charset_normalizer-3.4.3-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:1606f4a55c0fd363d754049cdf400175ee96c992b1f8018b993941f221221c5f", size = 151830 },
-    { url = "https://files.pythonhosted.org/packages/ae/02/e29e22b4e02839a0e4a06557b1999d0a47db3567e82989b5bb21f3fbbd9f/charset_normalizer-3.4.3-cp312-cp312-musllinux_1_2_aarch64.whl", hash = "sha256:027b776c26d38b7f15b26a5da1044f376455fb3766df8fc38563b4efbc515154", size = 148854 },
-    { url = "https://files.pythonhosted.org/packages/05/6b/e2539a0a4be302b481e8cafb5af8792da8093b486885a1ae4d15d452bcec/charset_normalizer-3.4.3-cp312-cp312-musllinux_1_2_ppc64le.whl", hash = "sha256:42e5088973e56e31e4fa58eb6bd709e42fc03799c11c42929592889a2e54c491", size = 160670 },
-    { url = "https://files.pythonhosted.org/packages/31/e7/883ee5676a2ef217a40ce0bffcc3d0dfbf9e64cbcfbdf822c52981c3304b/charset_normalizer-3.4.3-cp312-cp312-musllinux_1_2_s390x.whl", hash = "sha256:cc34f233c9e71701040d772aa7490318673aa7164a0efe3172b2981218c26d93", size = 158501 },
-    { url = "https://files.pythonhosted.org/packages/c1/35/6525b21aa0db614cf8b5792d232021dca3df7f90a1944db934efa5d20bb1/charset_normalizer-3.4.3-cp312-cp312-musllinux_1_2_x86_64.whl", hash = "sha256:320e8e66157cc4e247d9ddca8e21f427efc7a04bbd0ac8a9faf56583fa543f9f", size = 153173 },
-    { url = "https://files.pythonhosted.org/packages/50/ee/f4704bad8201de513fdc8aac1cabc87e38c5818c93857140e06e772b5892/charset_normalizer-3.4.3-cp312-cp312-win32.whl", hash = "sha256:fb6fecfd65564f208cbf0fba07f107fb661bcd1a7c389edbced3f7a493f70e37", size = 99822 },
-    { url = "https://files.pythonhosted.org/packages/39/f5/3b3836ca6064d0992c58c7561c6b6eee1b3892e9665d650c803bd5614522/charset_normalizer-3.4.3-cp312-cp312-win_amd64.whl", hash = "sha256:86df271bf921c2ee3818f0522e9a5b8092ca2ad8b065ece5d7d9d0e9f4849bcc", size = 107543 },
-    { url = "https://files.pythonhosted.org/packages/65/ca/2135ac97709b400c7654b4b764daf5c5567c2da45a30cdd20f9eefe2d658/charset_normalizer-3.4.3-cp313-cp313-macosx_10_13_universal2.whl", hash = "sha256:14c2a87c65b351109f6abfc424cab3927b3bdece6f706e4d12faaf3d52ee5efe", size = 205326 },
-    { url = "https://files.pythonhosted.org/packages/71/11/98a04c3c97dd34e49c7d247083af03645ca3730809a5509443f3c37f7c99/charset_normalizer-3.4.3-cp313-cp313-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:41d1fc408ff5fdfb910200ec0e74abc40387bccb3252f3f27c0676731df2b2c8", size = 146008 },
-    { url = "https://files.pythonhosted.org/packages/60/f5/4659a4cb3c4ec146bec80c32d8bb16033752574c20b1252ee842a95d1a1e/charset_normalizer-3.4.3-cp313-cp313-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:1bb60174149316da1c35fa5233681f7c0f9f514509b8e399ab70fea5f17e45c9", size = 159196 },
-    { url = "https://files.pythonhosted.org/packages/86/9e/f552f7a00611f168b9a5865a1414179b2c6de8235a4fa40189f6f79a1753/charset_normalizer-3.4.3-cp313-cp313-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:30d006f98569de3459c2fc1f2acde170b7b2bd265dc1943e87e1a4efe1b67c31", size = 156819 },
-    { url = "https://files.pythonhosted.org/packages/7e/95/42aa2156235cbc8fa61208aded06ef46111c4d3f0de233107b3f38631803/charset_normalizer-3.4.3-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:416175faf02e4b0810f1f38bcb54682878a4af94059a1cd63b8747244420801f", size = 151350 },
-    { url = "https://files.pythonhosted.org/packages/c2/a9/3865b02c56f300a6f94fc631ef54f0a8a29da74fb45a773dfd3dcd380af7/charset_normalizer-3.4.3-cp313-cp313-musllinux_1_2_aarch64.whl", hash = "sha256:6aab0f181c486f973bc7262a97f5aca3ee7e1437011ef0c2ec04b5a11d16c927", size = 148644 },
-    { url = "https://files.pythonhosted.org/packages/77/d9/cbcf1a2a5c7d7856f11e7ac2d782aec12bdfea60d104e60e0aa1c97849dc/charset_normalizer-3.4.3-cp313-cp313-musllinux_1_2_ppc64le.whl", hash = "sha256:fdabf8315679312cfa71302f9bd509ded4f2f263fb5b765cf1433b39106c3cc9", size = 160468 },
-    { url = "https://files.pythonhosted.org/packages/f6/42/6f45efee8697b89fda4d50580f292b8f7f9306cb2971d4b53f8914e4d890/charset_normalizer-3.4.3-cp313-cp313-musllinux_1_2_s390x.whl", hash = "sha256:bd28b817ea8c70215401f657edef3a8aa83c29d447fb0b622c35403780ba11d5", size = 158187 },
-    { url = "https://files.pythonhosted.org/packages/70/99/f1c3bdcfaa9c45b3ce96f70b14f070411366fa19549c1d4832c935d8e2c3/charset_normalizer-3.4.3-cp313-cp313-musllinux_1_2_x86_64.whl", hash = "sha256:18343b2d246dc6761a249ba1fb13f9ee9a2bcd95decc767319506056ea4ad4dc", size = 152699 },
-    { url = "https://files.pythonhosted.org/packages/a3/ad/b0081f2f99a4b194bcbb1934ef3b12aa4d9702ced80a37026b7607c72e58/charset_normalizer-3.4.3-cp313-cp313-win32.whl", hash = "sha256:6fb70de56f1859a3f71261cbe41005f56a7842cc348d3aeb26237560bfa5e0ce", size = 99580 },
-    { url = "https://files.pythonhosted.org/packages/9a/8f/ae790790c7b64f925e5c953b924aaa42a243fb778fed9e41f147b2a5715a/charset_normalizer-3.4.3-cp313-cp313-win_amd64.whl", hash = "sha256:cf1ebb7d78e1ad8ec2a8c4732c7be2e736f6e5123a4146c5b89c9d1f585f8cef", size = 107366 },
-    { url = "https://files.pythonhosted.org/packages/8e/91/b5a06ad970ddc7a0e513112d40113e834638f4ca1120eb727a249fb2715e/charset_normalizer-3.4.3-cp314-cp314-macosx_10_13_universal2.whl", hash = "sha256:3cd35b7e8aedeb9e34c41385fda4f73ba609e561faedfae0a9e75e44ac558a15", size = 204342 },
-    { url = "https://files.pythonhosted.org/packages/ce/ec/1edc30a377f0a02689342f214455c3f6c2fbedd896a1d2f856c002fc3062/charset_normalizer-3.4.3-cp314-cp314-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b89bc04de1d83006373429975f8ef9e7932534b8cc9ca582e4db7d20d91816db", size = 145995 },
-    { url = "https://files.pythonhosted.org/packages/17/e5/5e67ab85e6d22b04641acb5399c8684f4d37caf7558a53859f0283a650e9/charset_normalizer-3.4.3-cp314-cp314-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:2001a39612b241dae17b4687898843f254f8748b796a2e16f1051a17078d991d", size = 158640 },
-    { url = "https://files.pythonhosted.org/packages/f1/e5/38421987f6c697ee3722981289d554957c4be652f963d71c5e46a262e135/charset_normalizer-3.4.3-cp314-cp314-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:8dcfc373f888e4fb39a7bc57e93e3b845e7f462dacc008d9749568b1c4ece096", size = 156636 },
-    { url = "https://files.pythonhosted.org/packages/a0/e4/5a075de8daa3ec0745a9a3b54467e0c2967daaaf2cec04c845f73493e9a1/charset_normalizer-3.4.3-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:18b97b8404387b96cdbd30ad660f6407799126d26a39ca65729162fd810a99aa", size = 150939 },
-    { url = "https://files.pythonhosted.org/packages/02/f7/3611b32318b30974131db62b4043f335861d4d9b49adc6d57c1149cc49d4/charset_normalizer-3.4.3-cp314-cp314-musllinux_1_2_aarch64.whl", hash = "sha256:ccf600859c183d70eb47e05a44cd80a4ce77394d1ac0f79dbd2dd90a69a3a049", size = 148580 },
-    { url = "https://files.pythonhosted.org/packages/7e/61/19b36f4bd67f2793ab6a99b979b4e4f3d8fc754cbdffb805335df4337126/charset_normalizer-3.4.3-cp314-cp314-musllinux_1_2_ppc64le.whl", hash = "sha256:53cd68b185d98dde4ad8990e56a58dea83a4162161b1ea9272e5c9182ce415e0", size = 159870 },
-    { url = "https://files.pythonhosted.org/packages/06/57/84722eefdd338c04cf3030ada66889298eaedf3e7a30a624201e0cbe424a/charset_normalizer-3.4.3-cp314-cp314-musllinux_1_2_s390x.whl", hash = "sha256:30a96e1e1f865f78b030d65241c1ee850cdf422d869e9028e2fc1d5e4db73b92", size = 157797 },
-    { url = "https://files.pythonhosted.org/packages/72/2a/aff5dd112b2f14bcc3462c312dce5445806bfc8ab3a7328555da95330e4b/charset_normalizer-3.4.3-cp314-cp314-musllinux_1_2_x86_64.whl", hash = "sha256:d716a916938e03231e86e43782ca7878fb602a125a91e7acb8b5112e2e96ac16", size = 152224 },
-    { url = "https://files.pythonhosted.org/packages/b7/8c/9839225320046ed279c6e839d51f028342eb77c91c89b8ef2549f951f3ec/charset_normalizer-3.4.3-cp314-cp314-win32.whl", hash = "sha256:c6dbd0ccdda3a2ba7c2ecd9d77b37f3b5831687d8dc1b6ca5f56a4880cc7b7ce", size = 100086 },
-    { url = "https://files.pythonhosted.org/packages/ee/7a/36fbcf646e41f710ce0a563c1c9a343c6edf9be80786edeb15b6f62e17db/charset_normalizer-3.4.3-cp314-cp314-win_amd64.whl", hash = "sha256:73dc19b562516fc9bcf6e5d6e596df0b4eb98d87e4f79f3ae71840e6ed21361c", size = 107400 },
-    { url = "https://files.pythonhosted.org/packages/22/82/63a45bfc36f73efe46731a3a71cb84e2112f7e0b049507025ce477f0f052/charset_normalizer-3.4.3-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:0f2be7e0cf7754b9a30eb01f4295cc3d4358a479843b31f328afd210e2c7598c", size = 198805 },
-    { url = "https://files.pythonhosted.org/packages/0c/52/8b0c6c3e53f7e546a5e49b9edb876f379725914e1130297f3b423c7b71c5/charset_normalizer-3.4.3-cp38-cp38-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:c60e092517a73c632ec38e290eba714e9627abe9d301c8c8a12ec32c314a2a4b", size = 142862 },
-    { url = "https://files.pythonhosted.org/packages/59/c0/a74f3bd167d311365e7973990243f32c35e7a94e45103125275b9e6c479f/charset_normalizer-3.4.3-cp38-cp38-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:252098c8c7a873e17dd696ed98bbe91dbacd571da4b87df3736768efa7a792e4", size = 155104 },
-    { url = "https://files.pythonhosted.org/packages/1a/79/ae516e678d6e32df2e7e740a7be51dc80b700e2697cb70054a0f1ac2c955/charset_normalizer-3.4.3-cp38-cp38-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:3653fad4fe3ed447a596ae8638b437f827234f01a8cd801842e43f3d0a6b281b", size = 152598 },
-    { url = "https://files.pythonhosted.org/packages/00/bd/ef9c88464b126fa176f4ef4a317ad9b6f4d30b2cffbc43386062367c3e2c/charset_normalizer-3.4.3-cp38-cp38-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8999f965f922ae054125286faf9f11bc6932184b93011d138925a1773830bbe9", size = 147391 },
-    { url = "https://files.pythonhosted.org/packages/7a/03/cbb6fac9d3e57f7e07ce062712ee80d80a5ab46614684078461917426279/charset_normalizer-3.4.3-cp38-cp38-musllinux_1_2_aarch64.whl", hash = "sha256:d95bfb53c211b57198bb91c46dd5a2d8018b3af446583aab40074bf7988401cb", size = 145037 },
-    { url = "https://files.pythonhosted.org/packages/64/d1/f9d141c893ef5d4243bc75c130e95af8fd4bc355beff06e9b1e941daad6e/charset_normalizer-3.4.3-cp38-cp38-musllinux_1_2_ppc64le.whl", hash = "sha256:5b413b0b1bfd94dbf4023ad6945889f374cd24e3f62de58d6bb102c4d9ae534a", size = 156425 },
-    { url = "https://files.pythonhosted.org/packages/c5/35/9c99739250742375167bc1b1319cd1cec2bf67438a70d84b2e1ec4c9daa3/charset_normalizer-3.4.3-cp38-cp38-musllinux_1_2_s390x.whl", hash = "sha256:b5e3b2d152e74e100a9e9573837aba24aab611d39428ded46f4e4022ea7d1942", size = 153734 },
-    { url = "https://files.pythonhosted.org/packages/50/10/c117806094d2c956ba88958dab680574019abc0c02bcf57b32287afca544/charset_normalizer-3.4.3-cp38-cp38-musllinux_1_2_x86_64.whl", hash = "sha256:a2d08ac246bb48479170408d6c19f6385fa743e7157d716e144cad849b2dd94b", size = 148551 },
-    { url = "https://files.pythonhosted.org/packages/61/c5/dc3ba772489c453621ffc27e8978a98fe7e41a93e787e5e5bde797f1dddb/charset_normalizer-3.4.3-cp38-cp38-win32.whl", hash = "sha256:ec557499516fc90fd374bf2e32349a2887a876fbf162c160e3c01b6849eaf557", size = 98459 },
-    { url = "https://files.pythonhosted.org/packages/05/35/bb59b1cd012d7196fc81c2f5879113971efc226a63812c9cf7f89fe97c40/charset_normalizer-3.4.3-cp38-cp38-win_amd64.whl", hash = "sha256:5d8d01eac18c423815ed4f4a2ec3b439d654e55ee4ad610e153cf02faf67ea40", size = 105887 },
-    { url = "https://files.pythonhosted.org/packages/c2/ca/9a0983dd5c8e9733565cf3db4df2b0a2e9a82659fd8aa2a868ac6e4a991f/charset_normalizer-3.4.3-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:70bfc5f2c318afece2f5838ea5e4c3febada0be750fcf4775641052bbba14d05", size = 207520 },
-    { url = "https://files.pythonhosted.org/packages/39/c6/99271dc37243a4f925b09090493fb96c9333d7992c6187f5cfe5312008d2/charset_normalizer-3.4.3-cp39-cp39-manylinux2014_aarch64.manylinux_2_17_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:23b6b24d74478dc833444cbd927c338349d6ae852ba53a0d02a2de1fce45b96e", size = 147307 },
-    { url = "https://files.pythonhosted.org/packages/e4/69/132eab043356bba06eb333cc2cc60c6340857d0a2e4ca6dc2b51312886b3/charset_normalizer-3.4.3-cp39-cp39-manylinux2014_ppc64le.manylinux_2_17_ppc64le.manylinux_2_28_ppc64le.whl", hash = "sha256:34a7f768e3f985abdb42841e20e17b330ad3aaf4bb7e7aeeb73db2e70f077b99", size = 160448 },
-    { url = "https://files.pythonhosted.org/packages/04/9a/914d294daa4809c57667b77470533e65def9c0be1ef8b4c1183a99170e9d/charset_normalizer-3.4.3-cp39-cp39-manylinux2014_s390x.manylinux_2_17_s390x.manylinux_2_28_s390x.whl", hash = "sha256:fb731e5deb0c7ef82d698b0f4c5bb724633ee2a489401594c5c88b02e6cb15f7", size = 157758 },
-    { url = "https://files.pythonhosted.org/packages/b0/a8/6f5bcf1bcf63cb45625f7c5cadca026121ff8a6c8a3256d8d8cd59302663/charset_normalizer-3.4.3-cp39-cp39-manylinux2014_x86_64.manylinux_2_17_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:257f26fed7d7ff59921b78244f3cd93ed2af1800ff048c33f624c87475819dd7", size = 152487 },
-    { url = "https://files.pythonhosted.org/packages/c4/72/d3d0e9592f4e504f9dea08b8db270821c909558c353dc3b457ed2509f2fb/charset_normalizer-3.4.3-cp39-cp39-musllinux_1_2_aarch64.whl", hash = "sha256:1ef99f0456d3d46a50945c98de1774da86f8e992ab5c77865ea8b8195341fc19", size = 150054 },
-    { url = "https://files.pythonhosted.org/packages/20/30/5f64fe3981677fe63fa987b80e6c01042eb5ff653ff7cec1b7bd9268e54e/charset_normalizer-3.4.3-cp39-cp39-musllinux_1_2_ppc64le.whl", hash = "sha256:2c322db9c8c89009a990ef07c3bcc9f011a3269bc06782f916cd3d9eed7c9312", size = 161703 },
-    { url = "https://files.pythonhosted.org/packages/e1/ef/dd08b2cac9284fd59e70f7d97382c33a3d0a926e45b15fc21b3308324ffd/charset_normalizer-3.4.3-cp39-cp39-musllinux_1_2_s390x.whl", hash = "sha256:511729f456829ef86ac41ca78c63a5cb55240ed23b4b737faca0eb1abb1c41bc", size = 159096 },
-    { url = "https://files.pythonhosted.org/packages/45/8c/dcef87cfc2b3f002a6478f38906f9040302c68aebe21468090e39cde1445/charset_normalizer-3.4.3-cp39-cp39-musllinux_1_2_x86_64.whl", hash = "sha256:88ab34806dea0671532d3f82d82b85e8fc23d7b2dd12fa837978dad9bb392a34", size = 153852 },
-    { url = "https://files.pythonhosted.org/packages/63/86/9cbd533bd37883d467fcd1bd491b3547a3532d0fbb46de2b99feeebf185e/charset_normalizer-3.4.3-cp39-cp39-win32.whl", hash = "sha256:16a8770207946ac75703458e2c743631c79c59c5890c80011d536248f8eaa432", size = 99840 },
-    { url = "https://files.pythonhosted.org/packages/ce/d6/7e805c8e5c46ff9729c49950acc4ee0aeb55efb8b3a56687658ad10c3216/charset_normalizer-3.4.3-cp39-cp39-win_amd64.whl", hash = "sha256:d22dbedd33326a4a5190dd4fe9e9e693ef12160c77382d9e87919bce54f3d4ca", size = 107438 },
-    { url = "https://files.pythonhosted.org/packages/8a/1f/f041989e93b001bc4e44bb1669ccdcf54d3f00e628229a85b08d330615c5/charset_normalizer-3.4.3-py3-none-any.whl", hash = "sha256:ce571ab16d890d23b5c278547ba694193a45011ff86a9162a71307ed9f86759a", size = 53175 },
-]
-
-[[package]]
-name = "click"
-version = "8.1.8"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "colorama", marker = "platform_system == 'Windows'" },
-    { name = "importlib-metadata", marker = "python_full_version < '3.8'" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/b9/2e/0090cbf739cee7d23781ad4b89a9894a41538e4fcf4c31dcdd705b78eb8b/click-8.1.8.tar.gz", hash = "sha256:ed53c9d8990d83c2a27deae68e4ee337473f6330c040a31d4225c9574d16096a", size = 226593 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/7e/d4/7ebdbd03970677812aac39c869717059dbb71a4cfc033ca6e5221787892c/click-8.1.8-py3-none-any.whl", hash = "sha256:63c132bbbed01578a06712a2d1f497bb62d9c1c0d329b7903a866228027263b2", size = 98188 },
-]
-
-[[package]]
-name = "colorama"
-version = "0.4.6"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/d8/53/6f443c9a4a8358a93a6792e2acffb9d9d5cb0a5cfd8802644b7b1c9a02e4/colorama-0.4.6.tar.gz", hash = "sha256:08695f5cb7ed6e0531a20572697297273c47b8cae5a63ffc6d6ed5c201be6e44", size = 27697 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/d1/d6/3965ed04c63042e047cb6a3e6ed1a63a35087b6a609aa3a15ed8ac56c221/colorama-0.4.6-py2.py3-none-any.whl", hash = "sha256:4f1d9991f5acc0ca119f9d443620b77f9d6b33703e51011c16baf57afb285fc6", size = 25335 },
-]
-
-[[package]]
-name = "distlib"
-version = "0.4.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/96/8e/709914eb2b5749865801041647dc7f4e6d00b549cfe88b65ca192995f07c/distlib-0.4.0.tar.gz", hash = "sha256:feec40075be03a04501a973d81f633735b4b69f98b05450592310c0f401a4e0d", size = 614605 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/33/6b/e0547afaf41bf2c42e52430072fa5658766e3d65bd4b03a563d1b6336f57/distlib-0.4.0-py2.py3-none-any.whl", hash = "sha256:9659f7d87e46584a30b5780e43ac7a2143098441670ff0a49d5f9034c54a6c16", size = 469047 },
-]
-
-[[package]]
-name = "exceptiongroup"
-version = "1.3.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "typing-extensions", marker = "python_full_version < '3.13'" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/0b/9f/a65090624ecf468cdca03533906e7c69ed7588582240cfe7cc9e770b50eb/exceptiongroup-1.3.0.tar.gz", hash = "sha256:b241f5885f560bc56a59ee63ca4c6a8bfa46ae4ad651af316d4e81817bb9fd88", size = 29749 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/36/f4/c6e662dade71f56cd2f3735141b265c3c79293c109549c1e6933b0651ffc/exceptiongroup-1.3.0-py3-none-any.whl", hash = "sha256:4d111e6e0c13d0644cad6ddaa7ed0261a0b36971f6d23e7ec9b4b9097da78a10", size = 16674 },
-]
-
-[[package]]
-name = "filelock"
-version = "3.12.2"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/00/0b/c506e9e44e4c4b6c89fcecda23dc115bf8e7ff7eb127e0cb9c114cbc9a15/filelock-3.12.2.tar.gz", hash = "sha256:002740518d8aa59a26b0c76e10fb8c6e15eae825d34b6fdf670333fd7b938d81", size = 12441 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/00/45/ec3407adf6f6b5bf867a4462b2b0af27597a26bd3cd6e2534cb6ab029938/filelock-3.12.2-py3-none-any.whl", hash = "sha256:cbb791cdea2a72f23da6ac5b5269ab0a0d161e9ef0100e653b69049a7706d1ec", size = 10923 },
-]
-
-[[package]]
-name = "flask"
-version = "2.2.5"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "click" },
-    { name = "importlib-metadata", marker = "python_full_version < '3.10'" },
-    { name = "itsdangerous" },
-    { name = "jinja2" },
-    { name = "werkzeug" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/5f/76/a4d2c4436dda4b0a12c71e075c508ea7988a1066b06a575f6afe4fecc023/Flask-2.2.5.tar.gz", hash = "sha256:edee9b0a7ff26621bd5a8c10ff484ae28737a2410d99b0bb9a6850c7fb977aa0", size = 697814 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/9f/1a/8b6d48162861009d1e017a9740431c78d860809773b66cac220a11aa3310/Flask-2.2.5-py3-none-any.whl", hash = "sha256:58107ed83443e86067e41eff4631b058178191a355886f8e479e347fa1285fdf", size = 101817 },
-]
-
-[[package]]
-name = "flask-cloudflared"
-version = "0.0.14"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "flask" },
-    { name = "requests" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/49/cf/d01e9f44c042bb77194c83d00bc98bcf1cd8ff24cdf686191d8d57b8d4dc/flask_cloudflared-0.0.14.tar.gz", hash = "sha256:be88151714f306d042923c6996332de39519c24a1914692547d92789d098912a", size = 6157 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/f9/e0/d90aa747ae35adfcd6fa116936b1eb42a6e5d8c69419694862476cefd34b/flask_cloudflared-0.0.14-py3-none-any.whl", hash = "sha256:abd11c16030d5c68f91cf64efaecdebf09ab0dce10db43333aaebc790569bdbf", size = 6420 },
-]
-
-[[package]]
-name = "identify"
-version = "2.5.24"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/c4/f8/498e13e408d25ee6ff04aa0acbf91ad8e9caae74be91720fc0e811e649b7/identify-2.5.24.tar.gz", hash = "sha256:0aac67d5b4812498056d28a9a512a483f5085cc28640b02b258a59dac34301d4", size = 98886 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/4f/fd/2c46fba2bc032ba4c970bb8de59d25187087d7138a0ebf7c1dcc91d94f01/identify-2.5.24-py2.py3-none-any.whl", hash = "sha256:986dbfb38b1140e763e413e6feb44cd731faf72d1909543178aa79b0e258265d", size = 98826 },
-]
-
-[[package]]
-name = "idna"
-version = "3.10"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/f1/70/7703c29685631f5a7590aa73f1f1d3fa9a380e654b86af429e0934a32f7d/idna-3.10.tar.gz", hash = "sha256:12f65c9b470abda6dc35cf8e63cc574b1c52b11df2c86030af0ac09b01b13ea9", size = 190490 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/76/c6/c88e154df9c4e1a2a66ccf0005a88dfb2650c1dffb6f5ce603dfbd452ce3/idna-3.10-py3-none-any.whl", hash = "sha256:946d195a0d259cbba61165e88e65941f16e9b36ea6ddb97f00452bae8b1287d3", size = 70442 },
-]
-
-[[package]]
-name = "importlib-metadata"
-version = "6.7.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "typing-extensions", marker = "python_full_version < '3.8'" },
-    { name = "zipp" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/a3/82/f6e29c8d5c098b6be61460371c2c5591f4a335923639edec43b3830650a4/importlib_metadata-6.7.0.tar.gz", hash = "sha256:1aaf550d4f73e5d6783e7acb77aec43d49da8017410afae93822cc9cca98c4d4", size = 53569 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/ff/94/64287b38c7de4c90683630338cf28f129decbba0a44f0c6db35a873c73c4/importlib_metadata-6.7.0-py3-none-any.whl", hash = "sha256:cb52082e659e97afc5dac71e79de97d8681de3aa07ff18578330904a9d18e5b5", size = 22934 },
-]
-
-[[package]]
-name = "importlib-resources"
-version = "5.12.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "zipp", marker = "python_full_version < '3.10'" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/4e/a2/3cab1de83f95dd15297c15bdc04d50902391d707247cada1f021bbfe2149/importlib_resources-5.12.0.tar.gz", hash = "sha256:4be82589bf5c1d7999aedf2a45159d10cb3ca4f19b2271f8792bc8e6da7b22f6", size = 39894 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/38/71/c13ea695a4393639830bf96baea956538ba7a9d06fcce7cef10bfff20f72/importlib_resources-5.12.0-py3-none-any.whl", hash = "sha256:7b1deeebbf351c7578e09bf2f63fa2ce8b5ffec296e0d349139d43cca061a81a", size = 36211 },
-]
-
-[[package]]
-name = "iniconfig"
-version = "2.0.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/d7/4b/cbd8e699e64a6f16ca3a8220661b5f83792b3017d0f79807cb8708d33913/iniconfig-2.0.0.tar.gz", hash = "sha256:2d91e135bf72d31a410b17c16da610a82cb55f6b0477d1a902134b24a455b8b3", size = 4646 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/ef/a6/62565a6e1cf69e10f5727360368e451d4b7f58beeac6173dc9db836a5b46/iniconfig-2.0.0-py3-none-any.whl", hash = "sha256:b6a85871a79d2e3b22d2d1b94ac2824226a63c6b741c88f7ae975f18b6778374", size = 5892 },
-]
-
-[[package]]
-name = "itsdangerous"
-version = "2.1.2"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/7f/a1/d3fb83e7a61fa0c0d3d08ad0a94ddbeff3731c05212617dff3a94e097f08/itsdangerous-2.1.2.tar.gz", hash = "sha256:5dbbc68b317e5e42f327f9021763545dc3fc3bfe22e6deb96aaf1fc38874156a", size = 56143 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/68/5f/447e04e828f47465eeab35b5d408b7ebaaaee207f48b7136c5a7267a30ae/itsdangerous-2.1.2-py3-none-any.whl", hash = "sha256:2c2349112351b88699d8d4b6b075022c0808887cb7ad10069318a8b0bc88db44", size = 15749 },
-]
-
-[[package]]
-name = "jinja2"
-version = "3.1.6"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "markupsafe" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/df/bf/f7da0350254c0ed7c72f3e33cef02e048281fec7ecec5f032d4aac52226b/jinja2-3.1.6.tar.gz", hash = "sha256:0137fb05990d35f1275a587e9aee6d56da821fc83491a0fb838183be43f66d6d", size = 245115 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/62/a1/3d680cbfd5f4b8f15abc1d571870c5fc3e594bb582bc3b64ea099db13e56/jinja2-3.1.6-py3-none-any.whl", hash = "sha256:85ece4451f492d0c13c5dd7c13a64681a86afae63a5f347908daf103ce6d2f67", size = 134899 },
-]
-
-[[package]]
-name = "markupsafe"
-version = "2.1.5"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/87/5b/aae44c6655f3801e81aa3eef09dbbf012431987ba564d7231722f68df02d/MarkupSafe-2.1.5.tar.gz", hash = "sha256:d283d37a890ba4c1ae73ffadf8046435c76e7bc2247bbb63c00bd1a709c6544b", size = 19384 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/e4/54/ad5eb37bf9d51800010a74e4665425831a9db4e7c4e0fde4352e391e808e/MarkupSafe-2.1.5-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:a17a92de5231666cfbe003f0e4b9b3a7ae3afb1ec2845aadc2bacc93ff85febc", size = 18206 },
-    { url = "https://files.pythonhosted.org/packages/6a/4a/a4d49415e600bacae038c67f9fecc1d5433b9d3c71a4de6f33537b89654c/MarkupSafe-2.1.5-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:72b6be590cc35924b02c78ef34b467da4ba07e4e0f0454a2c5907f473fc50ce5", size = 14079 },
-    { url = "https://files.pythonhosted.org/packages/0a/7b/85681ae3c33c385b10ac0f8dd025c30af83c78cec1c37a6aa3b55e67f5ec/MarkupSafe-2.1.5-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:e61659ba32cf2cf1481e575d0462554625196a1f2fc06a1c777d3f48e8865d46", size = 26620 },
-    { url = "https://files.pythonhosted.org/packages/7c/52/2b1b570f6b8b803cef5ac28fdf78c0da318916c7d2fe9402a84d591b394c/MarkupSafe-2.1.5-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2174c595a0d73a3080ca3257b40096db99799265e1c27cc5a610743acd86d62f", size = 25818 },
-    { url = "https://files.pythonhosted.org/packages/29/fe/a36ba8c7ca55621620b2d7c585313efd10729e63ef81e4e61f52330da781/MarkupSafe-2.1.5-cp310-cp310-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ae2ad8ae6ebee9d2d94b17fb62763125f3f374c25618198f40cbb8b525411900", size = 25493 },
-    { url = "https://files.pythonhosted.org/packages/60/ae/9c60231cdfda003434e8bd27282b1f4e197ad5a710c14bee8bea8a9ca4f0/MarkupSafe-2.1.5-cp310-cp310-musllinux_1_1_aarch64.whl", hash = "sha256:075202fa5b72c86ad32dc7d0b56024ebdbcf2048c0ba09f1cde31bfdd57bcfff", size = 30630 },
-    { url = "https://files.pythonhosted.org/packages/65/dc/1510be4d179869f5dafe071aecb3f1f41b45d37c02329dfba01ff59e5ac5/MarkupSafe-2.1.5-cp310-cp310-musllinux_1_1_i686.whl", hash = "sha256:598e3276b64aff0e7b3451b72e94fa3c238d452e7ddcd893c3ab324717456bad", size = 29745 },
-    { url = "https://files.pythonhosted.org/packages/30/39/8d845dd7d0b0613d86e0ef89549bfb5f61ed781f59af45fc96496e897f3a/MarkupSafe-2.1.5-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:fce659a462a1be54d2ffcacea5e3ba2d74daa74f30f5f143fe0c58636e355fdd", size = 30021 },
-    { url = "https://files.pythonhosted.org/packages/c7/5c/356a6f62e4f3c5fbf2602b4771376af22a3b16efa74eb8716fb4e328e01e/MarkupSafe-2.1.5-cp310-cp310-win32.whl", hash = "sha256:d9fad5155d72433c921b782e58892377c44bd6252b5af2f67f16b194987338a4", size = 16659 },
-    { url = "https://files.pythonhosted.org/packages/69/48/acbf292615c65f0604a0c6fc402ce6d8c991276e16c80c46a8f758fbd30c/MarkupSafe-2.1.5-cp310-cp310-win_amd64.whl", hash = "sha256:bf50cd79a75d181c9181df03572cdce0fbb75cc353bc350712073108cba98de5", size = 17213 },
-    { url = "https://files.pythonhosted.org/packages/11/e7/291e55127bb2ae67c64d66cef01432b5933859dfb7d6949daa721b89d0b3/MarkupSafe-2.1.5-cp311-cp311-macosx_10_9_universal2.whl", hash = "sha256:629ddd2ca402ae6dbedfceeba9c46d5f7b2a61d9749597d4307f943ef198fc1f", size = 18219 },
-    { url = "https://files.pythonhosted.org/packages/6b/cb/aed7a284c00dfa7c0682d14df85ad4955a350a21d2e3b06d8240497359bf/MarkupSafe-2.1.5-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:5b7b716f97b52c5a14bffdf688f971b2d5ef4029127f1ad7a513973cfd818df2", size = 14098 },
-    { url = "https://files.pythonhosted.org/packages/1c/cf/35fe557e53709e93feb65575c93927942087e9b97213eabc3fe9d5b25a55/MarkupSafe-2.1.5-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:6ec585f69cec0aa07d945b20805be741395e28ac1627333b1c5b0105962ffced", size = 29014 },
-    { url = "https://files.pythonhosted.org/packages/97/18/c30da5e7a0e7f4603abfc6780574131221d9148f323752c2755d48abad30/MarkupSafe-2.1.5-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:b91c037585eba9095565a3556f611e3cbfaa42ca1e865f7b8015fe5c7336d5a5", size = 28220 },
-    { url = "https://files.pythonhosted.org/packages/0c/40/2e73e7d532d030b1e41180807a80d564eda53babaf04d65e15c1cf897e40/MarkupSafe-2.1.5-cp311-cp311-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:7502934a33b54030eaf1194c21c692a534196063db72176b0c4028e140f8f32c", size = 27756 },
-    { url = "https://files.pythonhosted.org/packages/18/46/5dca760547e8c59c5311b332f70605d24c99d1303dd9a6e1fc3ed0d73561/MarkupSafe-2.1.5-cp311-cp311-musllinux_1_1_aarch64.whl", hash = "sha256:0e397ac966fdf721b2c528cf028494e86172b4feba51d65f81ffd65c63798f3f", size = 33988 },
-    { url = "https://files.pythonhosted.org/packages/6d/c5/27febe918ac36397919cd4a67d5579cbbfa8da027fa1238af6285bb368ea/MarkupSafe-2.1.5-cp311-cp311-musllinux_1_1_i686.whl", hash = "sha256:c061bb86a71b42465156a3ee7bd58c8c2ceacdbeb95d05a99893e08b8467359a", size = 32718 },
-    { url = "https://files.pythonhosted.org/packages/f8/81/56e567126a2c2bc2684d6391332e357589a96a76cb9f8e5052d85cb0ead8/MarkupSafe-2.1.5-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:3a57fdd7ce31c7ff06cdfbf31dafa96cc533c21e443d57f5b1ecc6cdc668ec7f", size = 33317 },
-    { url = "https://files.pythonhosted.org/packages/00/0b/23f4b2470accb53285c613a3ab9ec19dc944eaf53592cb6d9e2af8aa24cc/MarkupSafe-2.1.5-cp311-cp311-win32.whl", hash = "sha256:397081c1a0bfb5124355710fe79478cdbeb39626492b15d399526ae53422b906", size = 16670 },
-    { url = "https://files.pythonhosted.org/packages/b7/a2/c78a06a9ec6d04b3445a949615c4c7ed86a0b2eb68e44e7541b9d57067cc/MarkupSafe-2.1.5-cp311-cp311-win_amd64.whl", hash = "sha256:2b7c57a4dfc4f16f7142221afe5ba4e093e09e728ca65c51f5620c9aaeb9a617", size = 17224 },
-    { url = "https://files.pythonhosted.org/packages/53/bd/583bf3e4c8d6a321938c13f49d44024dbe5ed63e0a7ba127e454a66da974/MarkupSafe-2.1.5-cp312-cp312-macosx_10_9_universal2.whl", hash = "sha256:8dec4936e9c3100156f8a2dc89c4b88d5c435175ff03413b443469c7c8c5f4d1", size = 18215 },
-    { url = "https://files.pythonhosted.org/packages/48/d6/e7cd795fc710292c3af3a06d80868ce4b02bfbbf370b7cee11d282815a2a/MarkupSafe-2.1.5-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:3c6b973f22eb18a789b1460b4b91bf04ae3f0c4234a0a6aa6b0a92f6f7b951d4", size = 14069 },
-    { url = "https://files.pythonhosted.org/packages/51/b5/5d8ec796e2a08fc814a2c7d2584b55f889a55cf17dd1a90f2beb70744e5c/MarkupSafe-2.1.5-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ac07bad82163452a6884fe8fa0963fb98c2346ba78d779ec06bd7a6262132aee", size = 29452 },
-    { url = "https://files.pythonhosted.org/packages/0a/0d/2454f072fae3b5a137c119abf15465d1771319dfe9e4acbb31722a0fff91/MarkupSafe-2.1.5-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:f5dfb42c4604dddc8e4305050aa6deb084540643ed5804d7455b5df8fe16f5e5", size = 28462 },
-    { url = "https://files.pythonhosted.org/packages/2d/75/fd6cb2e68780f72d47e6671840ca517bda5ef663d30ada7616b0462ad1e3/MarkupSafe-2.1.5-cp312-cp312-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ea3d8a3d18833cf4304cd2fc9cbb1efe188ca9b5efef2bdac7adc20594a0e46b", size = 27869 },
-    { url = "https://files.pythonhosted.org/packages/b0/81/147c477391c2750e8fc7705829f7351cf1cd3be64406edcf900dc633feb2/MarkupSafe-2.1.5-cp312-cp312-musllinux_1_1_aarch64.whl", hash = "sha256:d050b3361367a06d752db6ead6e7edeb0009be66bc3bae0ee9d97fb326badc2a", size = 33906 },
-    { url = "https://files.pythonhosted.org/packages/8b/ff/9a52b71839d7a256b563e85d11050e307121000dcebc97df120176b3ad93/MarkupSafe-2.1.5-cp312-cp312-musllinux_1_1_i686.whl", hash = "sha256:bec0a414d016ac1a18862a519e54b2fd0fc8bbfd6890376898a6c0891dd82e9f", size = 32296 },
-    { url = "https://files.pythonhosted.org/packages/88/07/2dc76aa51b481eb96a4c3198894f38b480490e834479611a4053fbf08623/MarkupSafe-2.1.5-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:58c98fee265677f63a4385256a6d7683ab1832f3ddd1e66fe948d5880c21a169", size = 33038 },
-    { url = "https://files.pythonhosted.org/packages/96/0c/620c1fb3661858c0e37eb3cbffd8c6f732a67cd97296f725789679801b31/MarkupSafe-2.1.5-cp312-cp312-win32.whl", hash = "sha256:8590b4ae07a35970728874632fed7bd57b26b0102df2d2b233b6d9d82f6c62ad", size = 16572 },
-    { url = "https://files.pythonhosted.org/packages/3f/14/c3554d512d5f9100a95e737502f4a2323a1959f6d0d01e0d0997b35f7b10/MarkupSafe-2.1.5-cp312-cp312-win_amd64.whl", hash = "sha256:823b65d8706e32ad2df51ed89496147a42a2a6e01c13cfb6ffb8b1e92bc910bb", size = 17127 },
-    { url = "https://files.pythonhosted.org/packages/a7/88/a940e11827ea1c136a34eca862486178294ae841164475b9ab216b80eb8e/MarkupSafe-2.1.5-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:c8b29db45f8fe46ad280a7294f5c3ec36dbac9491f2d1c17345be8e69cc5928f", size = 13982 },
-    { url = "https://files.pythonhosted.org/packages/cb/06/0d28bd178db529c5ac762a625c335a9168a7a23f280b4db9c95e97046145/MarkupSafe-2.1.5-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:ec6a563cff360b50eed26f13adc43e61bc0c04d94b8be985e6fb24b81f6dcfdf", size = 26335 },
-    { url = "https://files.pythonhosted.org/packages/4a/1d/c4f5016f87ced614eacc7d5fb85b25bcc0ff53e8f058d069fc8cbfdc3c7a/MarkupSafe-2.1.5-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a549b9c31bec33820e885335b451286e2969a2d9e24879f83fe904a5ce59d70a", size = 25557 },
-    { url = "https://files.pythonhosted.org/packages/b3/fb/c18b8c9fbe69e347fdbf782c6478f1bc77f19a830588daa224236678339b/MarkupSafe-2.1.5-cp37-cp37m-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:4f11aa001c540f62c6166c7726f71f7573b52c68c31f014c25cc7901deea0b52", size = 25245 },
-    { url = "https://files.pythonhosted.org/packages/2f/69/30d29adcf9d1d931c75001dd85001adad7374381c9c2086154d9f6445be6/MarkupSafe-2.1.5-cp37-cp37m-musllinux_1_1_aarch64.whl", hash = "sha256:7b2e5a267c855eea6b4283940daa6e88a285f5f2a67f2220203786dfa59b37e9", size = 31013 },
-    { url = "https://files.pythonhosted.org/packages/3a/03/63498d05bd54278b6ca340099e5b52ffb9cdf2ee4f2d9b98246337e21689/MarkupSafe-2.1.5-cp37-cp37m-musllinux_1_1_i686.whl", hash = "sha256:2d2d793e36e230fd32babe143b04cec8a8b3eb8a3122d2aceb4a371e6b09b8df", size = 30178 },
-    { url = "https://files.pythonhosted.org/packages/68/79/11b4fe15124692f8673b603433e47abca199a08ecd2a4851bfbdc97dc62d/MarkupSafe-2.1.5-cp37-cp37m-musllinux_1_1_x86_64.whl", hash = "sha256:ce409136744f6521e39fd8e2a24c53fa18ad67aa5bc7c2cf83645cce5b5c4e50", size = 30429 },
-    { url = "https://files.pythonhosted.org/packages/ed/88/408bdbf292eb86f03201c17489acafae8358ba4e120d92358308c15cea7c/MarkupSafe-2.1.5-cp37-cp37m-win32.whl", hash = "sha256:4096e9de5c6fdf43fb4f04c26fb114f61ef0bf2e5604b6ee3019d51b69e8c371", size = 16633 },
-    { url = "https://files.pythonhosted.org/packages/6c/4c/3577a52eea1880538c435176bc85e5b3379b7ab442327ccd82118550758f/MarkupSafe-2.1.5-cp37-cp37m-win_amd64.whl", hash = "sha256:4275d846e41ecefa46e2015117a9f491e57a71ddd59bbead77e904dc02b1bed2", size = 17215 },
-    { url = "https://files.pythonhosted.org/packages/f8/ff/2c942a82c35a49df5de3a630ce0a8456ac2969691b230e530ac12314364c/MarkupSafe-2.1.5-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:656f7526c69fac7f600bd1f400991cc282b417d17539a1b228617081106feb4a", size = 18192 },
-    { url = "https://files.pythonhosted.org/packages/4f/14/6f294b9c4f969d0c801a4615e221c1e084722ea6114ab2114189c5b8cbe0/MarkupSafe-2.1.5-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:97cafb1f3cbcd3fd2b6fbfb99ae11cdb14deea0736fc2b0952ee177f2b813a46", size = 14072 },
-    { url = "https://files.pythonhosted.org/packages/81/d4/fd74714ed30a1dedd0b82427c02fa4deec64f173831ec716da11c51a50aa/MarkupSafe-2.1.5-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1f3fbcb7ef1f16e48246f704ab79d79da8a46891e2da03f8783a5b6fa41a9532", size = 26928 },
-    { url = "https://files.pythonhosted.org/packages/c7/bd/50319665ce81bb10e90d1cf76f9e1aa269ea6f7fa30ab4521f14d122a3df/MarkupSafe-2.1.5-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:fa9db3f79de01457b03d4f01b34cf91bc0048eb2c3846ff26f66687c2f6d16ab", size = 26106 },
-    { url = "https://files.pythonhosted.org/packages/4c/6f/f2b0f675635b05f6afd5ea03c094557bdb8622fa8e673387444fe8d8e787/MarkupSafe-2.1.5-cp38-cp38-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:ffee1f21e5ef0d712f9033568f8344d5da8cc2869dbd08d87c84656e6a2d2f68", size = 25781 },
-    { url = "https://files.pythonhosted.org/packages/51/e0/393467cf899b34a9d3678e78961c2c8cdf49fb902a959ba54ece01273fb1/MarkupSafe-2.1.5-cp38-cp38-musllinux_1_1_aarch64.whl", hash = "sha256:5dedb4db619ba5a2787a94d877bc8ffc0566f92a01c0ef214865e54ecc9ee5e0", size = 30518 },
-    { url = "https://files.pythonhosted.org/packages/f6/02/5437e2ad33047290dafced9df741d9efc3e716b75583bbd73a9984f1b6f7/MarkupSafe-2.1.5-cp38-cp38-musllinux_1_1_i686.whl", hash = "sha256:30b600cf0a7ac9234b2638fbc0fb6158ba5bdcdf46aeb631ead21248b9affbc4", size = 29669 },
-    { url = "https://files.pythonhosted.org/packages/0e/7d/968284145ffd9d726183ed6237c77938c021abacde4e073020f920e060b2/MarkupSafe-2.1.5-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:8dd717634f5a044f860435c1d8c16a270ddf0ef8588d4887037c5028b859b0c3", size = 29933 },
-    { url = "https://files.pythonhosted.org/packages/bf/f3/ecb00fc8ab02b7beae8699f34db9357ae49d9f21d4d3de6f305f34fa949e/MarkupSafe-2.1.5-cp38-cp38-win32.whl", hash = "sha256:daa4ee5a243f0f20d528d939d06670a298dd39b1ad5f8a72a4275124a7819eff", size = 16656 },
-    { url = "https://files.pythonhosted.org/packages/92/21/357205f03514a49b293e214ac39de01fadd0970a6e05e4bf1ddd0ffd0881/MarkupSafe-2.1.5-cp38-cp38-win_amd64.whl", hash = "sha256:619bc166c4f2de5caa5a633b8b7326fbe98e0ccbfacabd87268a2b15ff73a029", size = 17206 },
-    { url = "https://files.pythonhosted.org/packages/0f/31/780bb297db036ba7b7bbede5e1d7f1e14d704ad4beb3ce53fb495d22bc62/MarkupSafe-2.1.5-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:7a68b554d356a91cce1236aa7682dc01df0edba8d043fd1ce607c49dd3c1edcf", size = 18193 },
-    { url = "https://files.pythonhosted.org/packages/6c/77/d77701bbef72892affe060cdacb7a2ed7fd68dae3b477a8642f15ad3b132/MarkupSafe-2.1.5-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:db0b55e0f3cc0be60c1f19efdde9a637c32740486004f20d1cff53c3c0ece4d2", size = 14073 },
-    { url = "https://files.pythonhosted.org/packages/d9/a7/1e558b4f78454c8a3a0199292d96159eb4d091f983bc35ef258314fe7269/MarkupSafe-2.1.5-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:3e53af139f8579a6d5f7b76549125f0d94d7e630761a2111bc431fd820e163b8", size = 26486 },
-    { url = "https://files.pythonhosted.org/packages/5f/5a/360da85076688755ea0cceb92472923086993e86b5613bbae9fbc14136b0/MarkupSafe-2.1.5-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:17b950fccb810b3293638215058e432159d2b71005c74371d784862b7e4683f3", size = 25685 },
-    { url = "https://files.pythonhosted.org/packages/6a/18/ae5a258e3401f9b8312f92b028c54d7026a97ec3ab20bfaddbdfa7d8cce8/MarkupSafe-2.1.5-cp39-cp39-manylinux_2_5_i686.manylinux1_i686.manylinux_2_17_i686.manylinux2014_i686.whl", hash = "sha256:4c31f53cdae6ecfa91a77820e8b151dba54ab528ba65dfd235c80b086d68a465", size = 25338 },
-    { url = "https://files.pythonhosted.org/packages/0b/cc/48206bd61c5b9d0129f4d75243b156929b04c94c09041321456fd06a876d/MarkupSafe-2.1.5-cp39-cp39-musllinux_1_1_aarch64.whl", hash = "sha256:bff1b4290a66b490a2f4719358c0cdcd9bafb6b8f061e45c7a2460866bf50c2e", size = 30439 },
-    { url = "https://files.pythonhosted.org/packages/d1/06/a41c112ab9ffdeeb5f77bc3e331fdadf97fa65e52e44ba31880f4e7f983c/MarkupSafe-2.1.5-cp39-cp39-musllinux_1_1_i686.whl", hash = "sha256:bc1667f8b83f48511b94671e0e441401371dfd0f0a795c7daa4a3cd1dde55bea", size = 29531 },
-    { url = "https://files.pythonhosted.org/packages/02/8c/ab9a463301a50dab04d5472e998acbd4080597abc048166ded5c7aa768c8/MarkupSafe-2.1.5-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:5049256f536511ee3f7e1b3f87d1d1209d327e818e6ae1365e8653d7e3abb6a6", size = 29823 },
-    { url = "https://files.pythonhosted.org/packages/bc/29/9bc18da763496b055d8e98ce476c8e718dcfd78157e17f555ce6dd7d0895/MarkupSafe-2.1.5-cp39-cp39-win32.whl", hash = "sha256:00e046b6dd71aa03a41079792f8473dc494d564611a8f89bbbd7cb93295ebdcf", size = 16658 },
-    { url = "https://files.pythonhosted.org/packages/f6/f8/4da07de16f10551ca1f640c92b5f316f9394088b183c6a57183df6de5ae4/MarkupSafe-2.1.5-cp39-cp39-win_amd64.whl", hash = "sha256:fa173ec60341d6bb97a89f5ea19c85c5643c1e7dedebc22f5181eb73573142c5", size = 17211 },
-]
-
-[[package]]
-name = "nodeenv"
-version = "1.9.1"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/43/16/fc88b08840de0e0a72a2f9d8c6bae36be573e475a6326ae854bcc549fc45/nodeenv-1.9.1.tar.gz", hash = "sha256:6ec12890a2dab7946721edbfbcd91f3319c6ccc9aec47be7c7e6b7011ee6645f", size = 47437 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/d2/1d/1b658dbd2b9fa9c4c9f32accbfc0205d532c8c6194dc0f2a4c0428e7128a/nodeenv-1.9.1-py2.py3-none-any.whl", hash = "sha256:ba11c9782d29c27c70ffbdda2d7415098754709be8a7056d79a737cd901155c9", size = 22314 },
-]
-
-[[package]]
-name = "numpy"
-version = "1.21.6"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/45/b7/de7b8e67f2232c26af57c205aaad29fe17754f793404f59c8a730c7a191a/numpy-1.21.6.zip", hash = "sha256:ecb55251139706669fdec2ff073c98ef8e9a84473e51e716211b41aa0f18e656", size = 10274544 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/ff/c6/05ae3c7f75b596e1bb3d78131c331eada9376a03d1af9801bd40e4675023/numpy-1.21.6-cp310-cp310-macosx_10_9_universal2.whl", hash = "sha256:8737609c3bbdd48e380d463134a35ffad3b22dc56295eff6f79fd85bd0eeeb25", size = 27203123 },
-    { url = "https://files.pythonhosted.org/packages/4a/72/a3379f83172f1431d7949138373e3a24beed68184c9362dab1b4d465be26/numpy-1.21.6-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:fdffbfb6832cd0b300995a2b08b8f6fa9f6e856d562800fea9182316d99c4e8e", size = 16965603 },
-    { url = "https://files.pythonhosted.org/packages/26/e7/4a6f579af8186372b03e8480e47df309520d91cfead8759b64dd5ac62688/numpy-1.21.6-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:3820724272f9913b597ccd13a467cc492a0da6b05df26ea09e78b171a0bb9da6", size = 12364595 },
-    { url = "https://files.pythonhosted.org/packages/57/ba/d8cbdfd507b541bb247beff24d9d7304ac8ffc379cf585701187d45d4512/numpy-1.21.6-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:f17e562de9edf691a42ddb1eb4a5541c20dd3f9e65b09ded2beb0799c0cf29bb", size = 13019278 },
-    { url = "https://files.pythonhosted.org/packages/6f/7b/036000a55680e6c7eb81502b0aa27ce0ed65d4d8805613909967d9f8baf6/numpy-1.21.6-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:5f30427731561ce75d7048ac254dbe47a2ba576229250fb60f0fb74db96501a1", size = 15906004 },
-    { url = "https://files.pythonhosted.org/packages/b0/77/ff8bbe56ff6cbbdbdb8a641c67cee61e29b2e8bfbb18732c2e1d2961fe4d/numpy-1.21.6-cp310-cp310-win32.whl", hash = "sha256:d4bf4d43077db55589ffc9009c0ba0a94fa4908b9586d6ccce2e0b164c86303c", size = 11706182 },
-    { url = "https://files.pythonhosted.org/packages/ec/03/93702ca9c4bd61791e46c80ff1f24943febb2317484cf7e8207688bbbd95/numpy-1.21.6-cp310-cp310-win_amd64.whl", hash = "sha256:d136337ae3cc69aa5e447e78d8e1514be8c3ec9b54264e680cf0b4bd9011574f", size = 14008913 },
-    { url = "https://files.pythonhosted.org/packages/32/dd/43d8b2b2ebf424f6555271a4c9f5b50dc3cc0aafa66c72b4d36863f71358/numpy-1.21.6-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:6aaf96c7f8cebc220cdfc03f1d5a31952f027dda050e5a703a0d1c396075e3e7", size = 16894122 },
-    { url = "https://files.pythonhosted.org/packages/06/78/b184f13f5461812a17a90b380d70a93fa3532460f0af9d72b0d93d8bc4ff/numpy-1.21.6-cp37-cp37m-manylinux_2_12_i686.manylinux2010_i686.whl", hash = "sha256:67c261d6c0a9981820c3a149d255a76918278a6b03b6a036800359aba1256d46", size = 13670697 },
-    { url = "https://files.pythonhosted.org/packages/6d/ad/ff3b21ebfe79a4d25b4a4f8e5cf9fd44a204adb6b33c09010f566f51027a/numpy-1.21.6-cp37-cp37m-manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:a6be4cb0ef3b8c9250c19cc122267263093eee7edd4e3fa75395dfda8c17a8e2", size = 15702369 },
-    { url = "https://files.pythonhosted.org/packages/b7/0d/86662f93102e42545cdf031da4fddf0ace9030ec67478932a628afc5973b/numpy-1.21.6-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:7c4068a8c44014b2d55f3c3f574c376b2494ca9cc73d2f1bd692382b6dffe3db", size = 12974170 },
-    { url = "https://files.pythonhosted.org/packages/cd/eb/f6f3258e7b0e0cc5c327778312bf4ee4978c8514aa28e97119ee206f6e60/numpy-1.21.6-cp37-cp37m-win32.whl", hash = "sha256:7c7e5fa88d9ff656e067876e4736379cc962d185d5cd808014a8a928d529ef4e", size = 11680505 },
-    { url = "https://files.pythonhosted.org/packages/97/9f/da37cc4a188a1d5d203d65ab28d6504e17594b5342e0c1dc5610ee6f4535/numpy-1.21.6-cp37-cp37m-win_amd64.whl", hash = "sha256:bcb238c9c96c00d3085b264e5c1a1207672577b93fa666c3b14a45240b14123a", size = 13969236 },
-    { url = "https://files.pythonhosted.org/packages/b5/e2/b2df1f664d644e690b40179fc0a07c163c6decf986c7adee8a85a094e8ce/numpy-1.21.6-cp38-cp38-macosx_10_9_universal2.whl", hash = "sha256:82691fda7c3f77c90e62da69ae60b5ac08e87e775b09813559f8901a88266552", size = 27127297 },
-    { url = "https://files.pythonhosted.org/packages/5b/d4/be63d2bed7d10f443dee42469623326b6bc51c9e5cd096ebb7227bca456f/numpy-1.21.6-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:643843bcc1c50526b3a71cd2ee561cf0d8773f062c8cbaf9ffac9fdf573f83ab", size = 16926225 },
-    { url = "https://files.pythonhosted.org/packages/0d/21/036363516c06737135ee58741e9c0af4899348ce3c5f5e04379240edd090/numpy-1.21.6-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:357768c2e4451ac241465157a3e929b265dfac85d9214074985b1786244f2ef3", size = 12329630 },
-    { url = "https://files.pythonhosted.org/packages/6a/52/a1dcf14b8e81d49c14112663290ee2ed545bd04988170138284a613bd926/numpy-1.21.6-cp38-cp38-manylinux_2_12_i686.manylinux2010_i686.whl", hash = "sha256:9f411b2c3f3d76bba0865b35a425157c5dcf54937f82bbeb3d3c180789dd66a6", size = 13719672 },
-    { url = "https://files.pythonhosted.org/packages/d5/43/e88bb1fb7d040ae8e0e06e749341b13f57701aab11fe9d71c99af6202c5c/numpy-1.21.6-cp38-cp38-manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:4aa48afdce4660b0076a00d80afa54e8a97cd49f457d68a4342d188a09451c1a", size = 15749422 },
-    { url = "https://files.pythonhosted.org/packages/86/c7/3f68d0a8dcc9458879c614707e6ffaf64a108664cfbba9702d3ba7ca4c82/numpy-1.21.6-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:d6a96eef20f639e6a97d23e57dd0c1b1069a7b4fd7027482a4c5c451cd7732f4", size = 13026205 },
-    { url = "https://files.pythonhosted.org/packages/6f/47/453023bd298f8b0be092d8a8bdd4b21f87a8c639ecb724a94cd75e23d216/numpy-1.21.6-cp38-cp38-win32.whl", hash = "sha256:5c3c8def4230e1b959671eb959083661b4a0d2e9af93ee339c7dada6759a9470", size = 11707551 },
-    { url = "https://files.pythonhosted.org/packages/48/5f/db4550e1c68206814a577ebd92c0dd082f3628fd7fc96725d44a521b0c92/numpy-1.21.6-cp38-cp38-win_amd64.whl", hash = "sha256:bf2ec4b75d0e9356edea834d1de42b31fe11f726a81dfb2c2112bc1eaa508fcf", size = 14009785 },
-    { url = "https://files.pythonhosted.org/packages/83/eb/a6a0d7fc8e718776c5c710692ea027607104710cba813c4b869182179334/numpy-1.21.6-cp39-cp39-macosx_10_9_universal2.whl", hash = "sha256:4391bd07606be175aafd267ef9bea87cf1b8210c787666ce82073b05f202add1", size = 27202394 },
-    { url = "https://files.pythonhosted.org/packages/4c/62/07402945bd5d5cf515a5f0cbc7263abf02ec0ddf3b19fbdc4af7537cd4d0/numpy-1.21.6-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:67f21981ba2f9d7ba9ade60c9e8cbaa8cf8e9ae51673934480e45cf55e953673", size = 16965163 },
-    { url = "https://files.pythonhosted.org/packages/44/56/041e886b4a8da813b7ec297c270fb3582d2ae8b7f33e106eb5c7a5e9184c/numpy-1.21.6-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:ee5ec40fdd06d62fe5d4084bef4fd50fd4bb6bfd2bf519365f569dc470163ab0", size = 12364236 },
-    { url = "https://files.pythonhosted.org/packages/61/f4/f01a8989e53a437ad660ab86c91514bec3d5067393e4a844b259f5a103de/numpy-1.21.6-cp39-cp39-manylinux_2_12_i686.manylinux2010_i686.whl", hash = "sha256:1dbe1c91269f880e364526649a52eff93ac30035507ae980d2fed33aaee633ac", size = 13721583 },
-    { url = "https://files.pythonhosted.org/packages/e7/f2/0bdcf2c40ef144cbbc9e0947eea831a145a98b0e5f8438fc09cf7fda0b35/numpy-1.21.6-cp39-cp39-manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:d9caa9d5e682102453d96a0ee10c7241b72859b01a941a397fd965f23b3e016b", size = 15734025 },
-    { url = "https://files.pythonhosted.org/packages/76/7f/830cf169eede1b855538f962e3a70c31755db6423652695b813ed04ff54e/numpy-1.21.6-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:58459d3bad03343ac4b1b42ed14d571b8743dc80ccbf27444f266729df1d6f5b", size = 13020515 },
-    { url = "https://files.pythonhosted.org/packages/1b/b5/7178d5a22427a9195ac69d6ec150415734f7a7a19d1142f82b89ead1dac4/numpy-1.21.6-cp39-cp39-win32.whl", hash = "sha256:7f5ae4f304257569ef3b948810816bc87c9146e8c446053539947eedeaa32786", size = 11706846 },
-    { url = "https://files.pythonhosted.org/packages/4d/04/bcd62448f2e772bc90a73ba21bacaa19817ae9905ae639969462862bd071/numpy-1.21.6-cp39-cp39-win_amd64.whl", hash = "sha256:e31f0bb5928b793169b87e3d1e070f2342b22d5245c755e2b81caa29756246c3", size = 14008965 },
-    { url = "https://files.pythonhosted.org/packages/2e/5a/6f3e280a10de48395053a559bfcb3b2221b74b57d062c1d6307fc965f549/numpy-1.21.6-pp37-pypy37_pp73-manylinux_2_12_x86_64.manylinux2010_x86_64.whl", hash = "sha256:dd1c8f6bd65d07d3810b90d02eba7997e32abbdf1277a481d698969e921a3be0", size = 15159527 },
-]
-
-[[package]]
-name = "packaging"
-version = "24.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/ee/b5/b43a27ac7472e1818c4bafd44430e69605baefe1f34440593e0332ec8b4d/packaging-24.0.tar.gz", hash = "sha256:eb82c5e3e56209074766e6885bb04b8c38a0c015d0a30036ebe7ece34c9989e9", size = 147882 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/49/df/1fceb2f8900f8639e278b056416d49134fb8d84c5942ffaa01ad34782422/packaging-24.0-py3-none-any.whl", hash = "sha256:2ddfb553fdf02fb784c234c7ba6ccc288296ceabec964ad2eae3777778130bc5", size = 53488 },
-]
-
-[[package]]
-name = "platformdirs"
-version = "4.0.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "typing-extensions", marker = "python_full_version < '3.8'" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/31/28/e40d24d2e2eb23135f8533ad33d582359c7825623b1e022f9d460def7c05/platformdirs-4.0.0.tar.gz", hash = "sha256:cb633b2bcf10c51af60beb0ab06d2f1d69064b43abf4c185ca6b28865f3f9731", size = 19914 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/31/16/70be3b725073035aa5fc3229321d06e22e73e3e09f6af78dcfdf16c7636c/platformdirs-4.0.0-py3-none-any.whl", hash = "sha256:118c954d7e949b35437270383a3f2531e99dd93cf7ce4dc8340d3356d30f173b", size = 17562 },
-]
-
-[[package]]
-name = "pluggy"
-version = "1.2.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "importlib-metadata", marker = "python_full_version < '3.8'" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/8a/42/8f2833655a29c4e9cb52ee8a2be04ceac61bcff4a680fb338cbd3d1e322d/pluggy-1.2.0.tar.gz", hash = "sha256:d12f0c4b579b15f5e054301bb226ee85eeeba08ffec228092f8defbaa3a4c4b3", size = 61613 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/51/32/4a79112b8b87b21450b066e102d6608907f4c885ed7b04c3fdb085d4d6ae/pluggy-1.2.0-py3-none-any.whl", hash = "sha256:c2fd55a7d7a3863cba1a013e4e2414658b1d07b6bc57b3919e0c63c9abb99849", size = 17695 },
-]
-
-[[package]]
-name = "pre-commit"
-version = "2.21.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "cfgv" },
-    { name = "identify" },
-    { name = "importlib-metadata", marker = "python_full_version < '3.8'" },
-    { name = "nodeenv" },
-    { name = "pyyaml" },
-    { name = "virtualenv" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/6b/00/1637ae945c6e10838ef5c41965f1c864e59301811bb203e979f335608e7c/pre_commit-2.21.0.tar.gz", hash = "sha256:31ef31af7e474a8d8995027fefdfcf509b5c913ff31f2015b4ec4beb26a6f658", size = 174966 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/a6/6b/6cfe3a8b351b54f4b6c6d2ad4286804e3367f628dce379c603d3b96635f4/pre_commit-2.21.0-py2.py3-none-any.whl", hash = "sha256:e2f91727039fc39a92f58a588a25b87f936de6567eed4f0e673e0507edc75bad", size = 201938 },
-]
-
-[[package]]
-name = "pyarrow"
-version = "12.0.1"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "numpy" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/c5/68/d3410e975bebbf5be00c1238d0418345d8ec5d88b7a6c102211a1c967edd/pyarrow-12.0.1.tar.gz", hash = "sha256:cce317fc96e5b71107bf1f9f184d5e54e2bd14bbf3f9a3d62819961f0af86fec", size = 1015259 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/13/2f/a42dbdf34528c70bbd5736a968631e3c8c2f911aea89f9c49f6f834e83b5/pyarrow-12.0.1-cp310-cp310-macosx_10_14_x86_64.whl", hash = "sha256:6d288029a94a9bb5407ceebdd7110ba398a00412c5b0155ee9813a40d246c5df", size = 24731979 },
-    { url = "https://files.pythonhosted.org/packages/64/05/76bcbea6903957c6467f99fcc6aaf07ac5ea675c02e75881719949801335/pyarrow-12.0.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:345e1828efdbd9aa4d4de7d5676778aba384a2c3add896d995b23d368e60e5af", size = 22646475 },
-    { url = "https://files.pythonhosted.org/packages/60/94/e56483c49ae2acee47af880ab4e0af7749811a0142a584d45543957ee1b3/pyarrow-12.0.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:8d6009fdf8986332b2169314da482baed47ac053311c8934ac6651e614deacd6", size = 36386496 },
-    { url = "https://files.pythonhosted.org/packages/0d/c8/886acfcce7cb2f7552f538d2b6deafd4841f3de42902943db15f1b42313d/pyarrow-12.0.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:2d3c4cbbf81e6dd23fe921bc91dc4619ea3b79bc58ef10bce0f49bdafb103daf", size = 38938438 },
-    { url = "https://files.pythonhosted.org/packages/13/68/1bc83fe2de87e2b785c503ef28293b56a4405d632f2b443cf00265b12d90/pyarrow-12.0.1-cp310-cp310-win_amd64.whl", hash = "sha256:cdacf515ec276709ac8042c7d9bd5be83b4f5f39c6c037a17a60d7ebfd92c890", size = 21465027 },
-    { url = "https://files.pythonhosted.org/packages/90/1e/fb0177d214a77198083156d750358c0a3ff696c96b329f443ad5513d25b6/pyarrow-12.0.1-cp311-cp311-macosx_10_14_x86_64.whl", hash = "sha256:749be7fd2ff260683f9cc739cb862fb11be376de965a2a8ccbf2693b098db6c7", size = 24687651 },
-    { url = "https://files.pythonhosted.org/packages/a7/ca/a34c5dd3393644865b82ac5df66e52311fd4ae2fc073f62b68b8538a0da4/pyarrow-12.0.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:6895b5fb74289d055c43db3af0de6e16b07586c45763cb5e558d38b86a91e3a7", size = 22601949 },
-    { url = "https://files.pythonhosted.org/packages/6a/ba/571de5dc75831b9a0f9e8d23823c1286b5c940588d4d8c87aab535779d53/pyarrow-12.0.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:1887bdae17ec3b4c046fcf19951e71b6a619f39fa674f9881216173566c8f718", size = 36384544 },
-    { url = "https://files.pythonhosted.org/packages/c5/52/19832487e6834164c523386a1b047dd5539fcbb876196b6f5619dfdab465/pyarrow-12.0.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:e2c9cb8eeabbadf5fcfc3d1ddea616c7ce893db2ce4dcef0ac13b099ad7ca082", size = 38930218 },
-    { url = "https://files.pythonhosted.org/packages/53/98/823deb5d3ff75386d5ab19b90c25b4a3aca01299bb53f217f2cfa954329c/pyarrow-12.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:ce4aebdf412bd0eeb800d8e47db854f9f9f7e2f5a0220440acf219ddfddd4f63", size = 21453804 },
-    { url = "https://files.pythonhosted.org/packages/8f/56/10fab8ea743b9bfd954d8648e715e1a947d7e131858d9670f83770626059/pyarrow-12.0.1-cp37-cp37m-macosx_10_14_x86_64.whl", hash = "sha256:e0d8730c7f6e893f6db5d5b86eda42c0a130842d101992b581e2138e4d5663d3", size = 24707735 },
-    { url = "https://files.pythonhosted.org/packages/cf/e2/94791e4cbb8cb16d5c99d016003746fab9d97f127342ba6b817bf639c767/pyarrow-12.0.1-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:43364daec02f69fec89d2315f7fbfbeec956e0d991cbbef471681bd77875c40f", size = 36561574 },
-    { url = "https://files.pythonhosted.org/packages/00/bd/4c03789f723337723670e8cf8935b18e170e7af3813ec38e71d2cfb0bd93/pyarrow-12.0.1-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:051f9f5ccf585f12d7de836e50965b3c235542cc896959320d9776ab93f3b33d", size = 39110664 },
-    { url = "https://files.pythonhosted.org/packages/0b/82/7f70296eb5167bc3bcee96a1460315af109affd7fef43e750c8ee4aac17c/pyarrow-12.0.1-cp37-cp37m-win_amd64.whl", hash = "sha256:be2757e9275875d2a9c6e6052ac7957fbbfc7bc7370e4a036a9b893e96fedaba", size = 21471413 },
-    { url = "https://files.pythonhosted.org/packages/75/a2/87fe24ab2c6efc6ad2335a2fc6bc33363fc70f67f18a3c18c494a4783aa2/pyarrow-12.0.1-cp38-cp38-macosx_10_14_x86_64.whl", hash = "sha256:cf812306d66f40f69e684300f7af5111c11f6e0d89d6b733e05a3de44961529d", size = 24719478 },
-    { url = "https://files.pythonhosted.org/packages/01/7c/427358d04bfcb217a165911bfada09e2d1e74fed04e40eb02c0c317ed2c4/pyarrow-12.0.1-cp38-cp38-macosx_11_0_arm64.whl", hash = "sha256:459a1c0ed2d68671188b2118c63bac91eaef6fc150c77ddd8a583e3c795737bf", size = 22638986 },
-    { url = "https://files.pythonhosted.org/packages/25/72/9afe7e2b61482ddc361c796857c19f69b2035ae20deebe0c1a54cb602b21/pyarrow-12.0.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:85e705e33eaf666bbe508a16fd5ba27ca061e177916b7a317ba5a51bee43384c", size = 36490338 },
-    { url = "https://files.pythonhosted.org/packages/e1/91/676b6ef5181fd0229ec35477eb94ff55fc5114ebab7a4669db311ddc9385/pyarrow-12.0.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:9120c3eb2b1f6f516a3b7a9714ed860882d9ef98c4b17edcdc91d95b7528db60", size = 39043549 },
-    { url = "https://files.pythonhosted.org/packages/af/cd/9674a609185bb9197b2dc25e8e61e1f4a2531f0754f43e4d3790f885707e/pyarrow-12.0.1-cp38-cp38-win_amd64.whl", hash = "sha256:c780f4dc40460015d80fcd6a6140de80b615349ed68ef9adb653fe351778c9b3", size = 21517590 },
-    { url = "https://files.pythonhosted.org/packages/25/33/8fa80189ea3ea7ac0b35b33e715de0466a0ec5064abb07a5b7ab5fe4f6fe/pyarrow-12.0.1-cp39-cp39-macosx_10_14_x86_64.whl", hash = "sha256:a3c63124fc26bf5f95f508f5d04e1ece8cc23a8b0af2a1e6ab2b1ec3fdc91b24", size = 24770513 },
-    { url = "https://files.pythonhosted.org/packages/a7/fd/a1488faf625a86b2ebf83bb977e48d9514785edfe438d4dbccf6e527bcc8/pyarrow-12.0.1-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:b13329f79fa4472324f8d32dc1b1216616d09bd1e77cfb13104dec5463632c36", size = 22676028 },
-    { url = "https://files.pythonhosted.org/packages/8b/14/dbda2f416906090824e5b58134ebef504065798bbcc98c929ce712be80ed/pyarrow-12.0.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:bb656150d3d12ec1396f6dde542db1675a95c0cc8366d507347b0beed96e87ca", size = 36430612 },
-    { url = "https://files.pythonhosted.org/packages/54/a2/5976df95323c4ca2b7baba31cb7a2a61a17461706043239d38a8e9dc281e/pyarrow-12.0.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6251e38470da97a5b2e00de5c6a049149f7b2bd62f12fa5dbb9ac674119ba71a", size = 38983701 },
-    { url = "https://files.pythonhosted.org/packages/dc/45/31441c988329afed625a791a7d78f1cf2fcb40dcc86a1d61e081287516a8/pyarrow-12.0.1-cp39-cp39-win_amd64.whl", hash = "sha256:3de26da901216149ce086920547dfff5cd22818c9eab67ebc41e863a5883bac7", size = 21511440 },
-]
-
-[[package]]
-name = "pycparser"
-version = "2.21"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/5e/0b/95d387f5f4433cb0f53ff7ad859bd2c6051051cebbb564f139a999ab46de/pycparser-2.21.tar.gz", hash = "sha256:e644fdec12f7872f86c58ff790da456218b10f863970249516d60a5eaca77206", size = 170877 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/62/d5/5f610ebe421e85889f2e55e33b7f9a6795bd982198517d912eb1c76e1a53/pycparser-2.21-py2.py3-none-any.whl", hash = "sha256:8ee45429555515e1f6b185e78100aea234072576aa43ab53aefcae078162fca9", size = 118697 },
-]
-
-[[package]]
-name = "pytest"
-version = "7.4.4"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "colorama", marker = "sys_platform == 'win32'" },
-    { name = "exceptiongroup", marker = "python_full_version < '3.11'" },
-    { name = "importlib-metadata", marker = "python_full_version < '3.8'" },
-    { name = "iniconfig" },
-    { name = "packaging" },
-    { name = "pluggy" },
-    { name = "tomli", marker = "python_full_version < '3.11'" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/80/1f/9d8e98e4133ffb16c90f3b405c43e38d3abb715bb5d7a63a5a684f7e46a3/pytest-7.4.4.tar.gz", hash = "sha256:2cf0005922c6ace4a3e2ec8b4080eb0d9753fdc93107415332f50ce9e7994280", size = 1357116 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/51/ff/f6e8b8f39e08547faece4bd80f89d5a8de68a38b2d179cc1c4490ffa3286/pytest-7.4.4-py3-none-any.whl", hash = "sha256:b090cdf5ed60bf4c45261be03239c2c1c22df034fbffe691abe93cd80cea01d8", size = 325287 },
-]
-
-[[package]]
-name = "pyyaml"
-version = "6.0.1"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/cd/e5/af35f7ea75cf72f2cd079c95ee16797de7cd71f29ea7c68ae5ce7be1eda0/PyYAML-6.0.1.tar.gz", hash = "sha256:bfdf460b1736c775f2ba9f6a92bca30bc2095067b8a9d77876d1fad6cc3b4a43", size = 125201 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/96/06/4beb652c0fe16834032e54f0956443d4cc797fe645527acee59e7deaa0a2/PyYAML-6.0.1-cp310-cp310-macosx_10_9_x86_64.whl", hash = "sha256:d858aa552c999bc8a8d57426ed01e40bef403cd8ccdd0fc5f6f04a00414cac2a", size = 189447 },
-    { url = "https://files.pythonhosted.org/packages/5b/07/10033a403b23405a8fc48975444463d3d10a5c2736b7eb2550b07b367429/PyYAML-6.0.1-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:fd66fc5d0da6d9815ba2cebeb4205f95818ff4b79c3ebe268e75d961704af52f", size = 169264 },
-    { url = "https://files.pythonhosted.org/packages/f1/26/55e4f21db1f72eaef092015d9017c11510e7e6301c62a6cfee91295d13c6/PyYAML-6.0.1-cp310-cp310-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:69b023b2b4daa7548bcfbd4aa3da05b3a74b772db9e23b982788168117739938", size = 677003 },
-    { url = "https://files.pythonhosted.org/packages/ba/91/090818dfa62e85181f3ae23dd1e8b7ea7f09684864a900cab72d29c57346/PyYAML-6.0.1-cp310-cp310-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:81e0b275a9ecc9c0c0c07b4b90ba548307583c125f54d5b6946cfee6360c733d", size = 699070 },
-    { url = "https://files.pythonhosted.org/packages/29/61/bf33c6c85c55bc45a29eee3195848ff2d518d84735eb0e2d8cb42e0d285e/PyYAML-6.0.1-cp310-cp310-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:ba336e390cd8e4d1739f42dfe9bb83a3cc2e80f567d8805e11b46f4a943f5515", size = 705525 },
-    { url = "https://files.pythonhosted.org/packages/07/91/45dfd0ef821a7f41d9d0136ea3608bb5b1653e42fd56a7970532cb5c003f/PyYAML-6.0.1-cp310-cp310-musllinux_1_1_x86_64.whl", hash = "sha256:326c013efe8048858a6d312ddd31d56e468118ad4cdeda36c719bf5bb6192290", size = 707514 },
-    { url = "https://files.pythonhosted.org/packages/b6/a0/b6700da5d49e9fed49dc3243d3771b598dad07abb37cc32e524607f96adc/PyYAML-6.0.1-cp310-cp310-win32.whl", hash = "sha256:bd4af7373a854424dabd882decdc5579653d7868b8fb26dc7d0e99f823aa5924", size = 130488 },
-    { url = "https://files.pythonhosted.org/packages/24/97/9b59b43431f98d01806b288532da38099cc6f2fea0f3d712e21e269c0279/PyYAML-6.0.1-cp310-cp310-win_amd64.whl", hash = "sha256:fd1592b3fdf65fff2ad0004b5e363300ef59ced41c2e6b3a99d4089fa8c5435d", size = 145338 },
-    { url = "https://files.pythonhosted.org/packages/ec/0d/26fb23e8863e0aeaac0c64e03fd27367ad2ae3f3cccf3798ee98ce160368/PyYAML-6.0.1-cp311-cp311-macosx_10_9_x86_64.whl", hash = "sha256:6965a7bc3cf88e5a1c3bd2e0b5c22f8d677dc88a455344035f03399034eb3007", size = 187867 },
-    { url = "https://files.pythonhosted.org/packages/28/09/55f715ddbf95a054b764b547f617e22f1d5e45d83905660e9a088078fe67/PyYAML-6.0.1-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:f003ed9ad21d6a4713f0a9b5a7a0a79e08dd0f221aff4525a2be4c346ee60aab", size = 167530 },
-    { url = "https://files.pythonhosted.org/packages/5e/94/7d5ee059dfb92ca9e62f4057dcdec9ac08a9e42679644854dc01177f8145/PyYAML-6.0.1-cp311-cp311-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:42f8152b8dbc4fe7d96729ec2b99c7097d656dc1213a3229ca5383f973a5ed6d", size = 732244 },
-    { url = "https://files.pythonhosted.org/packages/06/92/e0224aa6ebf9dc54a06a4609da37da40bb08d126f5535d81bff6b417b2ae/PyYAML-6.0.1-cp311-cp311-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:062582fca9fabdd2c8b54a3ef1c978d786e0f6b3a1510e0ac93ef59e0ddae2bc", size = 752871 },
-    { url = "https://files.pythonhosted.org/packages/7b/5e/efd033ab7199a0b2044dab3b9f7a4f6670e6a52c089de572e928d2873b06/PyYAML-6.0.1-cp311-cp311-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:d2b04aac4d386b172d5b9692e2d2da8de7bfb6c387fa4f801fbf6fb2e6ba4673", size = 757729 },
-    { url = "https://files.pythonhosted.org/packages/03/5c/c4671451b2f1d76ebe352c0945d4cd13500adb5d05f5a51ee296d80152f7/PyYAML-6.0.1-cp311-cp311-musllinux_1_1_x86_64.whl", hash = "sha256:e7d73685e87afe9f3b36c799222440d6cf362062f78be1013661b00c5c6f678b", size = 748528 },
-    { url = "https://files.pythonhosted.org/packages/73/9c/766e78d1efc0d1fca637a6b62cea1b4510a7fb93617eb805223294fef681/PyYAML-6.0.1-cp311-cp311-win32.whl", hash = "sha256:1635fd110e8d85d55237ab316b5b011de701ea0f29d07611174a1b42f1444741", size = 130286 },
-    { url = "https://files.pythonhosted.org/packages/b3/34/65bb4b2d7908044963ebf614fe0fdb080773fc7030d7e39c8d3eddcd4257/PyYAML-6.0.1-cp311-cp311-win_amd64.whl", hash = "sha256:bf07ee2fef7014951eeb99f56f39c9bb4af143d8aa3c21b1677805985307da34", size = 144699 },
-    { url = "https://files.pythonhosted.org/packages/bc/06/1b305bf6aa704343be85444c9d011f626c763abb40c0edc1cad13bfd7f86/PyYAML-6.0.1-cp312-cp312-macosx_10_9_x86_64.whl", hash = "sha256:855fb52b0dc35af121542a76b9a84f8d1cd886ea97c84703eaa6d88e37a2ad28", size = 178692 },
-    { url = "https://files.pythonhosted.org/packages/84/02/404de95ced348b73dd84f70e15a41843d817ff8c1744516bf78358f2ffd2/PyYAML-6.0.1-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:40df9b996c2b73138957fe23a16a4f0ba614f4c0efce1e9406a184b6d07fa3a9", size = 165622 },
-    { url = "https://files.pythonhosted.org/packages/c7/4c/4a2908632fc980da6d918b9de9c1d9d7d7e70b2672b1ad5166ed27841ef7/PyYAML-6.0.1-cp312-cp312-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a08c6f0fe150303c1c6b71ebcd7213c2858041a7e01975da3a99aed1e7a378ef", size = 696937 },
-    { url = "https://files.pythonhosted.org/packages/b4/33/720548182ffa8344418126017aa1d4ab4aeec9a2275f04ce3f3573d8ace8/PyYAML-6.0.1-cp312-cp312-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:6c22bec3fbe2524cde73d7ada88f6566758a8f7227bfbf93a408a9d86bcc12a0", size = 724969 },
-    { url = "https://files.pythonhosted.org/packages/4f/78/77b40157b6cb5f2d3d31a3d9b2efd1ba3505371f76730d267e8b32cf4b7f/PyYAML-6.0.1-cp312-cp312-musllinux_1_1_x86_64.whl", hash = "sha256:8d4e9c88387b0f5c7d5f281e55304de64cf7f9c0021a3525bd3b1c542da3b0e4", size = 712604 },
-    { url = "https://files.pythonhosted.org/packages/2e/97/3e0e089ee85e840f4b15bfa00e4e63d84a3691ababbfea92d6f820ea6f21/PyYAML-6.0.1-cp312-cp312-win32.whl", hash = "sha256:d483d2cdf104e7c9fa60c544d92981f12ad66a457afae824d146093b8c294c54", size = 126098 },
-    { url = "https://files.pythonhosted.org/packages/2b/9f/fbade56564ad486809c27b322d0f7e6a89c01f6b4fe208402e90d4443a99/PyYAML-6.0.1-cp312-cp312-win_amd64.whl", hash = "sha256:0d3304d8c0adc42be59c5f8a4d9e3d7379e6955ad754aa9d6ab7a398b59dd1df", size = 138675 },
-    { url = "https://files.pythonhosted.org/packages/c7/d1/02baa09d39b1bb1ebaf0d850d106d1bdcb47c91958557f471153c49dc03b/PyYAML-6.0.1-cp37-cp37m-macosx_10_9_x86_64.whl", hash = "sha256:b1275ad35a5d18c62a7220633c913e1b42d44b46ee12554e5fd39c70a243d6a3", size = 189627 },
-    { url = "https://files.pythonhosted.org/packages/e5/31/ba812efa640a264dbefd258986a5e4e786230cb1ee4a9f54eb28ca01e14a/PyYAML-6.0.1-cp37-cp37m-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:18aeb1bf9a78867dc38b259769503436b7c72f7a1f1f4c93ff9a17de54319b27", size = 658438 },
-    { url = "https://files.pythonhosted.org/packages/4d/f1/08f06159739254c8947899c9fc901241614195db15ba8802ff142237664c/PyYAML-6.0.1-cp37-cp37m-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:596106435fa6ad000c2991a98fa58eeb8656ef2325d7e158344fb33864ed87e3", size = 680304 },
-    { url = "https://files.pythonhosted.org/packages/d7/8f/db62b0df635b9008fe90aa68424e99cee05e68b398740c8a666a98455589/PyYAML-6.0.1-cp37-cp37m-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:baa90d3f661d43131ca170712d903e6295d1f7a0f595074f151c0aed377c9b9c", size = 670140 },
-    { url = "https://files.pythonhosted.org/packages/cc/5c/fcabd17918348c7db2eeeb0575705aaf3f7ab1657f6ce29b2e31737dd5d1/PyYAML-6.0.1-cp37-cp37m-win32.whl", hash = "sha256:9046c58c4395dff28dd494285c82ba00b546adfc7ef001486fbf0324bc174fba", size = 137577 },
-    { url = "https://files.pythonhosted.org/packages/1e/ae/964ccb88a938f20ece5754878f182cfbd846924930d02d29d06af8d4c69e/PyYAML-6.0.1-cp37-cp37m-win_amd64.whl", hash = "sha256:4fb147e7a67ef577a588a0e2c17b6db51dda102c71de36f8549b6816a96e1867", size = 153248 },
-    { url = "https://files.pythonhosted.org/packages/7f/5d/2779ea035ba1e533c32ed4a249b4e0448f583ba10830b21a3cddafe11a4e/PyYAML-6.0.1-cp38-cp38-macosx_10_9_x86_64.whl", hash = "sha256:1d4c7e777c441b20e32f52bd377e0c409713e8bb1386e1099c2415f26e479595", size = 191734 },
-    { url = "https://files.pythonhosted.org/packages/e1/a1/27bfac14b90adaaccf8c8289f441e9f76d94795ec1e7a8f134d9f2cb3d0b/PyYAML-6.0.1-cp38-cp38-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:a0cd17c15d3bb3fa06978b4e8958dcdc6e0174ccea823003a106c7d4d7899ac5", size = 723767 },
-    { url = "https://files.pythonhosted.org/packages/c1/39/47ed4d65beec9ce07267b014be85ed9c204fa373515355d3efa62d19d892/PyYAML-6.0.1-cp38-cp38-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:28c119d996beec18c05208a8bd78cbe4007878c6dd15091efb73a30e90539696", size = 749067 },
-    { url = "https://files.pythonhosted.org/packages/c8/6b/6600ac24725c7388255b2f5add93f91e58a5d7efaf4af244fdbcc11a541b/PyYAML-6.0.1-cp38-cp38-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:7e07cbde391ba96ab58e532ff4803f79c4129397514e1413a7dc761ccd755735", size = 736569 },
-    { url = "https://files.pythonhosted.org/packages/0d/46/62ae77677e532c0af6c81ddd6f3dbc16bdcc1208b077457354442d220bfb/PyYAML-6.0.1-cp38-cp38-musllinux_1_1_x86_64.whl", hash = "sha256:49a183be227561de579b4a36efbb21b3eab9651dd81b1858589f796549873dd6", size = 787738 },
-    { url = "https://files.pythonhosted.org/packages/d6/6a/439d1a6f834b9a9db16332ce16c4a96dd0e3970b65fe08cbecd1711eeb77/PyYAML-6.0.1-cp38-cp38-win32.whl", hash = "sha256:184c5108a2aca3c5b3d3bf9395d50893a7ab82a38004c8f61c258d4428e80206", size = 139797 },
-    { url = "https://files.pythonhosted.org/packages/29/0f/9782fa5b10152abf033aec56a601177ead85ee03b57781f2d9fced09eefc/PyYAML-6.0.1-cp38-cp38-win_amd64.whl", hash = "sha256:1e2722cc9fbb45d9b87631ac70924c11d3a401b2d7f410cc0e3bbf249f2dca62", size = 157350 },
-    { url = "https://files.pythonhosted.org/packages/57/c5/5d09b66b41d549914802f482a2118d925d876dc2a35b2d127694c1345c34/PyYAML-6.0.1-cp39-cp39-macosx_10_9_x86_64.whl", hash = "sha256:9eb6caa9a297fc2c2fb8862bc5370d0303ddba53ba97e71f08023b6cd73d16a8", size = 197846 },
-    { url = "https://files.pythonhosted.org/packages/0e/88/21b2f16cb2123c1e9375f2c93486e35fdc86e63f02e274f0e99c589ef153/PyYAML-6.0.1-cp39-cp39-macosx_11_0_arm64.whl", hash = "sha256:c8098ddcc2a85b61647b2590f825f3db38891662cfc2fc776415143f599bb859", size = 174396 },
-    { url = "https://files.pythonhosted.org/packages/ac/6c/967d91a8edf98d2b2b01d149bd9e51b8f9fb527c98d80ebb60c6b21d60c4/PyYAML-6.0.1-cp39-cp39-manylinux_2_17_aarch64.manylinux2014_aarch64.whl", hash = "sha256:5773183b6446b2c99bb77e77595dd486303b4faab2b086e7b17bc6bef28865f6", size = 731824 },
-    { url = "https://files.pythonhosted.org/packages/4a/4b/c71ef18ef83c82f99e6da8332910692af78ea32bd1d1d76c9787dfa36aea/PyYAML-6.0.1-cp39-cp39-manylinux_2_17_s390x.manylinux2014_s390x.whl", hash = "sha256:b786eecbdf8499b9ca1d697215862083bd6d2a99965554781d0d8d1ad31e13a0", size = 754777 },
-    { url = "https://files.pythonhosted.org/packages/7d/39/472f2554a0f1e825bd7c5afc11c817cd7a2f3657460f7159f691fbb37c51/PyYAML-6.0.1-cp39-cp39-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:bc1bf2925a1ecd43da378f4db9e4f799775d6367bdb94671027b73b393a7c42c", size = 738883 },
-    { url = "https://files.pythonhosted.org/packages/40/da/a175a35cf5583580e90ac3e2a3dbca90e43011593ae62ce63f79d7b28d92/PyYAML-6.0.1-cp39-cp39-musllinux_1_1_x86_64.whl", hash = "sha256:04ac92ad1925b2cff1db0cfebffb6ffc43457495c9b3c39d3fcae417d7125dc5", size = 750294 },
-    { url = "https://files.pythonhosted.org/packages/24/62/7fcc372442ec8ea331da18c24b13710e010c5073ab851ef36bf9dacb283f/PyYAML-6.0.1-cp39-cp39-win32.whl", hash = "sha256:faca3bdcf85b2fc05d06ff3fbc1f83e1391b3e724afa3feba7d13eeab355484c", size = 136936 },
-    { url = "https://files.pythonhosted.org/packages/84/4d/82704d1ab9290b03da94e6425f5e87396b999fd7eb8e08f3a92c158402bf/PyYAML-6.0.1-cp39-cp39-win_amd64.whl", hash = "sha256:510c9deebc5c0225e8c96813043e62b680ba2f9c50a08d3724c7f28a747d1486", size = 152751 },
-]
-
-[[package]]
-name = "requests"
-version = "2.31.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "certifi" },
-    { name = "charset-normalizer" },
-    { name = "idna" },
-    { name = "urllib3" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/9d/be/10918a2eac4ae9f02f6cfe6414b7a155ccd8f7f9d4380d62fd5b955065c3/requests-2.31.0.tar.gz", hash = "sha256:942c5a758f98d790eaed1a29cb6eefc7ffb0d1cf7af05c3d2791656dbd6ad1e1", size = 110794 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/70/8e/0e2d847013cb52cd35b38c009bb167a1a26b2ce6cd6965bf26b47bc0bf44/requests-2.31.0-py3-none-any.whl", hash = "sha256:58cd2187c01e70e6e26505bca751777aa9f2ee0b7f4300988b709f44e013003f", size = 62574 },
-]
-
-[[package]]
-name = "setuptools"
-version = "68.0.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/dc/98/5f896af066c128669229ff1aa81553ac14cfb3e5e74b6b44594132b8540e/setuptools-68.0.0.tar.gz", hash = "sha256:baf1fdb41c6da4cd2eae722e135500da913332ab3f2f5c7d33af9b492acb5235", size = 2194111 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/c7/42/be1c7bbdd83e1bfb160c94b9cafd8e25efc7400346cf7ccdbdb452c467fa/setuptools-68.0.0-py3-none-any.whl", hash = "sha256:11e52c67415a381d10d6b462ced9cfb97066179f0e871399e006c4ab101fc85f", size = 804037 },
-]
-
-[[package]]
-name = "six"
-version = "1.17.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/94/e7/b2c673351809dca68a0e064b6af791aa332cf192da575fd474ed7d6f16a2/six-1.17.0.tar.gz", hash = "sha256:ff70335d468e7eb6ec65b95b99d3a2836546063f63acc5171de367e834932a81", size = 34031 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/b7/ce/149a00dd41f10bc29e5921b496af8b574d8413afcd5e30dfa0ed46c2cc5e/six-1.17.0-py2.py3-none-any.whl", hash = "sha256:4721f391ed90541fddacab5acf947aa0d3dc7d27b2e1e8eda2be8970586c3274", size = 11050 },
-]
-
-[[package]]
-name = "tomli"
-version = "2.0.1"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/c0/3f/d7af728f075fb08564c5949a9c95e44352e23dee646869fa104a3b2060a3/tomli-2.0.1.tar.gz", hash = "sha256:de526c12914f0c550d15924c62d72abc48d6fe7364aa87328337a31007fe8a4f", size = 15164 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/97/75/10a9ebee3fd790d20926a90a2547f0bf78f371b2f13aa822c759680ca7b9/tomli-2.0.1-py3-none-any.whl", hash = "sha256:939de3e7a6161af0c887ef91b7d41a53e7c5a1ca976325f429cb46ea9bc30ecc", size = 12757 },
-]
-
-[[package]]
-name = "tqdm"
-version = "4.67.1"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "colorama", marker = "platform_system == 'Windows'" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/a8/4b/29b4ef32e036bb34e4ab51796dd745cdba7ed47ad142a9f4a1eb8e0c744d/tqdm-4.67.1.tar.gz", hash = "sha256:f8aef9c52c08c13a65f30ea34f4e5aac3fd1a34959879d7e59e63027286627f2", size = 169737 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/d0/30/dc54f88dd4a2b5dc8a0279bdd7270e735851848b762aeb1c1184ed1f6b14/tqdm-4.67.1-py3-none-any.whl", hash = "sha256:26445eca388f82e72884e0d580d5464cd801a3ea01e63e5601bdff9ba6a48de2", size = 78540 },
-]
-
-[[package]]
-name = "triton"
-version = "3.3.1"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "setuptools" },
-]
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/8d/a9/549e51e9b1b2c9b854fd761a1d23df0ba2fbc60bd0c13b489ffa518cfcb7/triton-3.3.1-cp310-cp310-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b74db445b1c562844d3cfad6e9679c72e93fdfb1a90a24052b03bb5c49d1242e", size = 155600257 },
-    { url = "https://files.pythonhosted.org/packages/21/2f/3e56ea7b58f80ff68899b1dbe810ff257c9d177d288c6b0f55bf2fe4eb50/triton-3.3.1-cp311-cp311-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b31e3aa26f8cb3cc5bf4e187bf737cbacf17311e1112b781d4a059353dfd731b", size = 155689937 },
-    { url = "https://files.pythonhosted.org/packages/24/5f/950fb373bf9c01ad4eb5a8cd5eaf32cdf9e238c02f9293557a2129b9c4ac/triton-3.3.1-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:9999e83aba21e1a78c1f36f21bce621b77bcaa530277a50484a7cb4a822f6e43", size = 155669138 },
-    { url = "https://files.pythonhosted.org/packages/74/1f/dfb531f90a2d367d914adfee771babbd3f1a5b26c3f5fbc458dee21daa78/triton-3.3.1-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:b89d846b5a4198317fec27a5d3a609ea96b6d557ff44b56c23176546023c4240", size = 155673035 },
-    { url = "https://files.pythonhosted.org/packages/28/71/bd20ffcb7a64c753dc2463489a61bf69d531f308e390ad06390268c4ea04/triton-3.3.1-cp313-cp313t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:a3198adb9d78b77818a5388bff89fa72ff36f9da0bc689db2f0a651a67ce6a42", size = 155735832 },
-    { url = "https://files.pythonhosted.org/packages/6d/81/ac4d50af22f594c4cb7c84fd2ad5ba1e0c03e2a83fe3483ddd79edcd7ec7/triton-3.3.1-cp39-cp39-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f6139aeb04a146b0b8e0fbbd89ad1e65861c57cfed881f21d62d3cb94a36bab7", size = 155596799 },
-]
-
-[[package]]
-name = "triton-viz"
-version = "2.0"
-source = { editable = "." }
-dependencies = [
-    { name = "anytree" },
-    { name = "cairocffi" },
-    { name = "flask" },
-    { name = "flask-cloudflared" },
-    { name = "pre-commit" },
-    { name = "pyarrow" },
-    { name = "pytest" },
-    { name = "setuptools" },
-    { name = "tqdm" },
-    { name = "triton" },
-    { name = "z3-solver" },
-]
-
-[package.metadata]
-requires-dist = [
-    { name = "anytree" },
-    { name = "cairocffi" },
-    { name = "flask" },
-    { name = "flask-cloudflared" },
-    { name = "pre-commit" },
-    { name = "pyarrow" },
-    { name = "pytest" },
-    { name = "setuptools" },
-    { name = "tqdm" },
-    { name = "triton" },
-    { name = "z3-solver" },
-]
-
-[[package]]
-name = "typing-extensions"
-version = "4.7.1"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/3c/8b/0111dd7d6c1478bf83baa1cab85c686426c7a6274119aceb2bd9d35395ad/typing_extensions-4.7.1.tar.gz", hash = "sha256:b75ddc264f0ba5615db7ba217daeb99701ad295353c45f9e95963337ceeeffb2", size = 72876 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/ec/6b/63cc3df74987c36fe26157ee12e09e8f9db4de771e0f3404263117e75b95/typing_extensions-4.7.1-py3-none-any.whl", hash = "sha256:440d5dd3af93b060174bf433bccd69b0babc3b15b1a8dca43789fd7f61514b36", size = 33232 },
-]
-
-[[package]]
-name = "urllib3"
-version = "2.0.7"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/af/47/b215df9f71b4fdba1025fc05a77db2ad243fa0926755a52c5e71659f4e3c/urllib3-2.0.7.tar.gz", hash = "sha256:c97dfde1f7bd43a71c8d2a58e369e9b2bf692d1334ea9f9cae55add7d0dd0f84", size = 282546 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/d2/b2/b157855192a68541a91ba7b2bbcb91f1b4faa51f8bae38d8005c034be524/urllib3-2.0.7-py3-none-any.whl", hash = "sha256:fdb6d215c776278489906c2f8916e6e7d4f5a9b602ccbcfdf7f016fc8da0596e", size = 124213 },
-]
-
-[[package]]
-name = "virtualenv"
-version = "20.26.6"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "distlib" },
-    { name = "filelock" },
-    { name = "importlib-metadata", marker = "python_full_version < '3.8'" },
-    { name = "platformdirs" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/3f/40/abc5a766da6b0b2457f819feab8e9203cbeae29327bd241359f866a3da9d/virtualenv-20.26.6.tar.gz", hash = "sha256:280aede09a2a5c317e409a00102e7077c6432c5a38f0ef938e643805a7ad2c48", size = 9372482 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/59/90/57b8ac0c8a231545adc7698c64c5a36fa7cd8e376c691b9bde877269f2eb/virtualenv-20.26.6-py3-none-any.whl", hash = "sha256:7345cc5b25405607a624d8418154577459c3e0277f5466dd79c49d5e492995f2", size = 5999862 },
-]
-
-[[package]]
-name = "werkzeug"
-version = "2.2.3"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "markupsafe" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/02/3c/baaebf3235c87d61d6593467056d5a8fba7c75ac838b8d100a5e64eba7a0/Werkzeug-2.2.3.tar.gz", hash = "sha256:2e1ccc9417d4da358b9de6f174e3ac094391ea1d4fbef2d667865d819dfd0afe", size = 845884 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/f6/f8/9da63c1617ae2a1dec2fbf6412f3a0cfe9d4ce029eccbda6e1e4258ca45f/Werkzeug-2.2.3-py3-none-any.whl", hash = "sha256:56433961bc1f12533306c624f3be5e744389ac61d722175d543e1751285da612", size = 233551 },
-]
-
-[[package]]
-name = "z3-solver"
-version = "4.15.3.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "importlib-resources", marker = "python_full_version < '3.9'" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/a3/60/9a924ee28cd1d12f2482834581d9024bf05110aa1098c056e847f05f7f76/z3_solver-4.15.3.0.tar.gz", hash = "sha256:78f69aebda5519bfd8af146a129f36cf4721a3c2667e80d9fe35cc9bb4d214a6", size = 4985945 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/63/45/dd8e9d7500faa05eafa589cc8f0f0b982ce575d51b455d62dab8e19dd571/z3_solver-4.15.3.0-py3-none-macosx_13_0_arm64.whl", hash = "sha256:65335aab295ded7c0ce27c85556067087a87052389ff160777d1a1d48ef0d74f", size = 36882388 },
-    { url = "https://files.pythonhosted.org/packages/54/9e/a11186061d9fead8be43bad7c75055585694124b2ccdd896ef249fe5824f/z3_solver-4.15.3.0-py3-none-macosx_13_0_x86_64.whl", hash = "sha256:3e62e93adff2def3537ff1ca67c3d58a6ca6d1944e0b5e774f88627b199d50e7", size = 39637842 },
-    { url = "https://files.pythonhosted.org/packages/b9/0b/f15168475e5493ea44fa3c5e642903f05d2b870db71ad05662ed87a06976/z3_solver-4.15.3.0-py3-none-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:a9afd9ceb290482097474d43f08415bcc1874f433189d1449f6c1508e9c68384", size = 29056003 },
-    { url = "https://files.pythonhosted.org/packages/a7/04/32a97b1f04175ec56213168ee3659709e876e3feacacb891ed3c26c1a82c/z3_solver-4.15.3.0-py3-none-manylinux_2_34_aarch64.whl", hash = "sha256:f61ef44552489077eedd7e6d9bed52ef1875decf86d66027742099a2703b1c77", size = 27074143 },
-    { url = "https://files.pythonhosted.org/packages/75/77/da54076a584557ea34f20800c68f725fe61f1dd987493fcb410b4a26f99f/z3_solver-4.15.3.0-py3-none-win32.whl", hash = "sha256:0c603f6bad7423d6411adda6af55030b725e3d30f54ea91b714abcedd73b848a", size = 13123666 },
-    { url = "https://files.pythonhosted.org/packages/c1/59/abc1bad8b25e9c576484ba65ca5ed225c8ed24d601ec242712f1c370b693/z3_solver-4.15.3.0-py3-none-win_amd64.whl", hash = "sha256:06abdf6c36f97c463aea827533504fd59476d015a65cf170a88bd6a53ba13ab5", size = 16203065 },
-]
-
-[[package]]
-name = "zipp"
-version = "3.15.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/00/27/f0ac6b846684cecce1ee93d32450c45ab607f65c2e0255f0092032d91f07/zipp-3.15.0.tar.gz", hash = "sha256:112929ad649da941c23de50f356a2b5570c954b65150642bccdd66bf194d224b", size = 18454 }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/5b/fa/c9e82bbe1af6266adf08afb563905eb87cab83fde00a0a08963510621047/zipp-3.15.0-py3-none-any.whl", hash = "sha256:48904fc76a60e542af151aded95726c1a5c34ed43ab4134b597665c86d7ad556", size = 6758 },
-]

From 9e2d4958bf0739de18c8e9d2f11e2f04bcc27379 Mon Sep 17 00:00:00 2001
From: gujialiang123 <562960306@qq.com>
Date: Fri, 17 Oct 2025 03:47:52 -0400
Subject: [PATCH 44/90] feat(flip): add Flip op tracing, 3D flip view, hover
 value API; fix example flip semantics; ui polish

---
 examples/flip.py                    |  81 ++++++++++++
 scripts/print_tracebacks.py         |  58 +++++++++
 triton_viz/clients/tracer/tracer.py |  25 +++-
 triton_viz/core/data.py             |  11 ++
 triton_viz/core/patch.py            |  93 ++++++++++++++
 triton_viz/static/flip.js           | 175 ++++++++++++++++++++++++++
 triton_viz/static/flip_3d.js        | 189 ++++++++++++++++++++++++++++
 triton_viz/static/flip_demo.js      | 178 ++++++++++++++++++++++++++
 triton_viz/static/gridblock.js      |  30 ++++-
 triton_viz/static/load.js           |   9 +-
 triton_viz/static/matmul.js         | 101 ++++++++-------
 triton_viz/static/store-utils.js    |   5 +-
 triton_viz/static/store.js          |   9 +-
 triton_viz/static/visualization.js  |   3 +-
 triton_viz/visualizer/draw.py       |  35 ++++++
 triton_viz/visualizer/interface.py  |  88 ++++++++++++-
 16 files changed, 1030 insertions(+), 60 deletions(-)
 create mode 100644 examples/flip.py
 create mode 100644 scripts/print_tracebacks.py
 create mode 100644 triton_viz/static/flip.js
 create mode 100644 triton_viz/static/flip_3d.js
 create mode 100644 triton_viz/static/flip_demo.js

diff --git a/examples/flip.py b/examples/flip.py
new file mode 100644
index 00000000..289250aa
--- /dev/null
+++ b/examples/flip.py
@@ -0,0 +1,81 @@
+import torch
+import triton
+import triton.language as tl
+
+import triton_viz
+from triton_viz.clients import Tracer
+
+
+@triton_viz.trace(clients=Tracer())
+@triton.jit
+def flip_1d_kernel(x_ptr, y_ptr, n, BLOCK: tl.constexpr):
+    pid = tl.program_id(0)
+    offs = pid * BLOCK + tl.arange(0, BLOCK)
+    mask = offs < n
+    x = tl.load(x_ptr + offs, mask=mask)
+    _ = tl.flip(x, dim=0)
+    rev_offs = (n - 1) - offs
+    tl.store(y_ptr + rev_offs, x, mask=mask)
+
+
+@triton_viz.trace(clients=Tracer())
+@triton.jit
+def flip_2d_kernel(
+    x_ptr,
+    y_ptr,
+    H,
+    W,
+    stride_h,
+    stride_w,
+    dim: tl.constexpr,
+    BLOCK_M: tl.constexpr,
+    BLOCK_N: tl.constexpr,
+):
+    pid_m = tl.program_id(0)
+    pid_n = tl.program_id(1)
+
+    offs_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+    offs_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)
+
+    mm = offs_m[:, None]
+    nn = offs_n[None, :]
+    mask = (mm < H) & (nn < W)
+
+    ptrs = x_ptr + mm * stride_h + nn * stride_w
+    vals = tl.load(ptrs, mask=mask, other=0)
+
+    _ = tl.flip(vals, dim=(0 if dim == 0 else 1))
+
+    if dim == 0:
+        rm = (H - 1) - mm
+        rn = nn
+    else:
+        rm = mm
+        rn = (W - 1) - nn
+
+    out_ptrs = y_ptr + rm * stride_h + rn * stride_w
+    tl.store(out_ptrs, vals, mask=mask)
+
+
+def run_1d(n: int = 256):
+    x = torch.arange(n, dtype=torch.int32)
+    y = torch.empty_like(x)
+    grid = (triton.cdiv(n, 128),)
+    flip_1d_kernel[grid](x, y, n, BLOCK=128)
+    assert torch.equal(y, torch.flip(x, dims=[0]))
+
+
+def run_2d(h: int = 16, w: int = 32, dim: int = 1):
+    x = torch.arange(h * w, dtype=torch.int32).reshape(h, w)
+    y = torch.empty_like(x)
+    grid = (triton.cdiv(h, 16), triton.cdiv(w, 16))
+    flip_2d_kernel[grid](
+        x, y, h, w, x.stride(0), x.stride(1), dim, BLOCK_M=16, BLOCK_N=16
+    )
+    assert torch.equal(y, torch.flip(x, dims=[dim]))
+
+
+if __name__ == "__main__":
+    run_1d(256)
+    run_2d(16, 32, dim=1)
+    triton_viz.launch(share=True, port=8003)
diff --git a/scripts/print_tracebacks.py b/scripts/print_tracebacks.py
new file mode 100644
index 00000000..f6686eca
--- /dev/null
+++ b/scripts/print_tracebacks.py
@@ -0,0 +1,58 @@
+import os
+import sys
+import torch
+import triton
+from triton_viz.core import config as cfg
+from triton_viz.core.trace import launches
+from triton_viz.core.data import Load, Store, Dot
+
+# Ensure project root is on sys.path so we can import examples
+PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
+if PROJECT_ROOT not in sys.path:
+    sys.path.insert(0, PROJECT_ROOT)
+
+from examples.load_store import simple_kernel
+
+
+def main():
+    # Reset viz config to a clean state
+    cfg.reset()
+
+    # Prepare sample data and run the kernel once
+    device = "cpu"
+    size = 16
+    BLOCK_SIZE = 8
+    torch.manual_seed(0)
+    x = torch.arange(size, dtype=torch.float32, device=device)
+    output = torch.empty_like(x)
+    grid = lambda meta: (triton.cdiv(size, meta["BLOCK_SIZE"]),)
+
+    # Run the traced kernel (decorated in examples/load_store.py)
+    simple_kernel[grid](x, output, size, BLOCK_SIZE)
+
+    # Inspect the last launch and print a sample op's call_path
+    print(f"Number of launches: {len(launches)}")
+    if not launches:
+        return
+
+    launch = launches[-1]
+    print(f"Number of records: {len(launch.records)}")
+
+    target = None
+    for rec in launch.records:
+        if isinstance(rec, (Load, Store, Dot)):
+            target = rec
+            break
+
+    if target is None:
+        print("No Load/Store/Dot record found to display call_path.")
+        return
+
+    print(f"=== {type(target).__name__} call_path ===")
+    for i, f in enumerate(target.call_path):
+        line = f.line if hasattr(f, "line") else ""
+        print(f"{i:02d}  {f.filename}:{f.lineno}  [{f.name}]  |  {line}")
+
+
+if __name__ == "__main__":
+    main()
diff --git a/triton_viz/clients/tracer/tracer.py b/triton_viz/clients/tracer/tracer.py
index 323cd96d..8ae55237 100644
--- a/triton_viz/clients/tracer/tracer.py
+++ b/triton_viz/clients/tracer/tracer.py
@@ -1,6 +1,6 @@
 from ...core.client import Client
 from ...core.callbacks import OpCallbacks, ForLoopCallbacks
-from ...core.data import Op, Load, Store, ReduceSum, Dot, Grid, RawLoad, RawStore
+from ...core.data import Op, Load, Store, ReduceSum, Dot, Grid, RawLoad, RawStore, Flip
 from typing import Callable, Optional, Union
 import numpy as np
 import traceback
@@ -155,6 +155,27 @@ def post_dot_callback(ret, input, other, *args):
             rec.call_path = _extract_user_frames()
             self.records.append(rec)
 
+        def post_flip_callback(ret, x, *args, **kwargs):
+            if not self.sample:
+                return
+            # Try to capture dim argument
+            dim = None
+            if args:
+                dim = args[0]
+            if "dim" in kwargs:
+                dim = kwargs.get("dim")
+            try:
+                in_shape = tuple(x.data.shape)
+                out_shape = tuple(ret.data.shape)
+            except Exception:
+                in_shape = getattr(getattr(x, "handle", None), "data", None)
+                out_shape = getattr(getattr(ret, "handle", None), "data", None)
+                in_shape = tuple(getattr(in_shape, "shape", []) or [])
+                out_shape = tuple(getattr(out_shape, "shape", []) or [])
+            rec = Flip(in_shape, out_shape, int(dim) if dim is not None else 0)
+            rec.call_path = _extract_user_frames()
+            self.records.append(rec)
+
         if op_type is Load:
             return OpCallbacks(before_callback=pre_load_callback)
         elif op_type is Store:
@@ -167,6 +188,8 @@ def post_dot_callback(ret, input, other, *args):
             return OpCallbacks(after_callback=post_reduce_sum_callback)
         elif op_type is Dot:
             return OpCallbacks(after_callback=post_dot_callback)
+        # Flip is wrapped at tl.flip; we don't have an interpreter op to hook here.
+        # The wrapper in patch_lang will append Flip records directly to tracer.
 
         return OpCallbacks()
 
diff --git a/triton_viz/core/data.py b/triton_viz/core/data.py
index 7f259e58..a4163583 100644
--- a/triton_viz/core/data.py
+++ b/triton_viz/core/data.py
@@ -104,6 +104,17 @@ def update_intermediate(self, row: int, col: int, result: float):
         self.intermediate_results[(row, col)] = result
 
 
+@dataclass
+class Flip(Op):
+    name: ClassVar[str] = "flip"
+    input_shape: tuple
+    output_shape: tuple
+    dim: int
+    # Optional payloads to help frontend render actual values when available
+    input_data: list | None = None
+    output_data: list | None = None
+
+
 @dataclass
 class MakeRange(Op):
     name: ClassVar[str] = "make_range"
diff --git a/triton_viz/core/patch.py b/triton_viz/core/patch.py
index af7288a6..975bfa4a 100644
--- a/triton_viz/core/patch.py
+++ b/triton_viz/core/patch.py
@@ -43,6 +43,7 @@
     Bitcast,
     AtomicCas,
 )
+from .data import Flip  # separate import to avoid reordering noise
 import inspect
 import ast
 from triton.runtime.interpreter import (
@@ -428,6 +429,95 @@ def unpatch_for_loop():
 def patch_lang(fn):
     triton_patch_lang(fn)
     fn.__globals__["_triton_viz_loop_patcher"] = _loop_patcher
+    # Wrap tl.flip to emit a Flip record after computing result
+    try:
+        _orig_flip = getattr(tl, "flip", None)
+        if _orig_flip is not None and not getattr(
+            _orig_flip, "__triton_viz_wrapped__", False
+        ):
+
+            def _viz_flip(x, *args, **kwargs):
+                # Call original flip implementation
+                ret = _orig_flip(x, *args, **kwargs)
+                # Best-effort extract dim
+                dim = None
+                if args:
+                    dim = args[0]
+                if "dim" in kwargs:
+                    dim = kwargs.get("dim")
+                # Best-effort shapes
+                in_shape = None
+                out_shape = None
+                x_arr = None
+                r_arr = None
+                try:
+                    # interpreter tensors may expose .data or .handle.data
+                    x_data = getattr(x, "data", None)
+                    if x_data is None and hasattr(x, "handle"):
+                        x_data = getattr(x.handle, "data", None)
+                    if x_data is not None:
+                        in_shape = tuple(x_data.shape)
+                        x_arr = x_data
+                except Exception:
+                    pass
+                try:
+                    r_data = getattr(ret, "data", None)
+                    if r_data is None and hasattr(ret, "handle"):
+                        r_data = getattr(ret.handle, "data", None)
+                    if r_data is not None:
+                        out_shape = tuple(r_data.shape)
+                        r_arr = r_data
+                except Exception:
+                    pass
+
+                # Emit a Flip record to the active tracer, if available
+                try:
+                    global _current_client_manager
+                    cm = _current_client_manager
+                    if cm is not None and hasattr(cm, "clients"):
+                        tracer = cm.get_client("tracer")
+                        if tracer is not None:
+                            input_payload = None
+                            output_payload = None
+                            try:
+                                # Avoid huge payloads: cap to 64k elements
+                                def _maybe_list(arr):
+                                    import numpy as _np
+
+                                    if arr is None:
+                                        return None
+                                    try:
+                                        if arr.size <= 65536:
+                                            return _np.asarray(arr).tolist()
+                                    except Exception:
+                                        pass
+                                    return None
+
+                                input_payload = _maybe_list(x_arr)
+                                output_payload = _maybe_list(r_arr)
+                            except Exception:
+                                pass
+
+                            rec = Flip(
+                                input_shape=in_shape or tuple(),
+                                output_shape=out_shape or (in_shape or tuple()),
+                                dim=int(dim) if dim is not None else 0,
+                                input_data=input_payload,
+                                output_data=output_payload,
+                            )
+                            # attach call path already handled by Flip.__post_init__
+                            tracer.records.append(rec)
+                except Exception:
+                    # Never fail kernel execution due to viz
+                    pass
+                return ret
+
+            # mark wrapper to avoid double-wrapping on subsequent patch_lang calls
+            setattr(_viz_flip, "__triton_viz_wrapped__", True)
+            tl.flip = _viz_flip  # type: ignore[assignment]
+    except Exception:
+        # If wrapping fails, continue without Flip records
+        pass
 
 
 def unpatch_lang():
@@ -478,6 +568,9 @@ def run_grid_loops(grid):
         k: v for k, v in kwargs.items() if k in argspec.args or k in triton_viz_args
     }
     client_manager = kwargs.pop("client_manager")
+    # Expose client_manager to tl.flip wrapper via a module-global
+    global _current_client_manager
+    _current_client_manager = client_manager
     args_hst, kwargs_hst = self._init_args_hst(args_dev, kwargs)
     # Prepare call arguments
     args = inspect.getcallargs(self.fn, *args_hst, **kwargs_hst)
diff --git a/triton_viz/static/flip.js b/triton_viz/static/flip.js
new file mode 100644
index 00000000..73f5758c
--- /dev/null
+++ b/triton_viz/static/flip.js
@@ -0,0 +1,175 @@
+import * as THREE from 'https://esm.sh/three@0.155.0/build/three.module.js';
+import { OrbitControls } from 'https://esm.sh/three@0.155.0/examples/jsm/controls/OrbitControls.js';
+import { setupScene, setupGeometries, createCube, CUBE_SIZE, GAP } from './load_utils.js';
+import { createFlip3D } from './flip_3d.js';
+
+export function createFlipVisualization(containerElement, op) {
+    const API_BASE = window.__TRITON_VIZ_API__ || '';
+    const overlay = document.createElement('div');
+    Object.assign(overlay.style, {
+        position: 'absolute', left: '0', top: '0', right: '0', bottom: '0',
+        zIndex: 3000, pointerEvents: 'auto'
+    });
+    containerElement.appendChild(overlay);
+
+    const { scene, camera, renderer } = setupScene(overlay, 0x15151b);
+    const controls = new OrbitControls(camera, renderer.domElement);
+    controls.enableDamping = true; controls.dampingFactor = 0.06;
+    const { cubeGeometry, edgesGeometry, lineMaterial } = setupGeometries();
+
+    // Helper color mapping
+    function hsv(h,s,v){
+        const c=v*s; const x=c*(1-Math.abs((h/60)%2-1)); const m=v-c; let r,g,b;
+        if(h<60){[r,g,b]=[c,x,0]} else if(h<120){[r,g,b]=[x,c,0]} else if(h<180){[r,g,b]=[0,c,x]}
+        else if(h<240){[r,g,b]=[0,x,c]} else if(h<300){[r,g,b]=[x,0,c]} else {[r,g,b]=[c,0,x]}
+        return new THREE.Color(r+m, g+m, b+m);
+    }
+    const shape = op.output_shape || op.input_shape || [];
+    const dim = op.dim || 0;
+    const is2D = (shape.length === 2);
+    const H = is2D ? (shape[0]||8) : 1;
+    const W = is2D ? (shape[1]||16) : ((shape[0]||16));
+
+    // Build input and output layers
+    const root = new THREE.Group(); scene.add(root);
+
+    // Input row
+    const inputGroup = new THREE.Group();
+    for (let r=0;r<H;r++){
+        for (let c=0;c<W;c++){
+            const idx = r*W + c;
+            const hue = (idx/((H*W)-1||1))*350;
+            const cube = createCube(hsv(hue,1,1), 'Input', c, r, 0, cubeGeometry, edgesGeometry, lineMaterial);
+            cube.position.set(c*(CUBE_SIZE+GAP), -r*(CUBE_SIZE+GAP), 0);
+            inputGroup.add(cube);
+        }
+    }
+    root.add(inputGroup);
+
+    // Output row (flipped mapping)
+    const outputGroup = new THREE.Group();
+    const mapCoord = (r, c) => {
+        if (!is2D) return { rr: 0, cc: (W-1-c) };
+        if (dim === 0) return { rr: (H-1-r), cc: c };
+        return { rr: r, cc: (W-1-c) };
+    };
+    for (let r=0;r<H;r++){
+        for (let c=0;c<W;c++){
+            const idx = r*W + c; const hue = (idx/((H*W)-1||1))*350;
+            const { rr, cc } = mapCoord(r, c);
+            const cube = createCube(hsv(hue,1,1), 'Output', cc, rr, 0, cubeGeometry, edgesGeometry, lineMaterial);
+            cube.position.set(cc*(CUBE_SIZE+GAP), -rr*(CUBE_SIZE+GAP), - (CUBE_SIZE*6));
+            outputGroup.add(cube);
+        }
+    }
+    root.add(outputGroup);
+
+    // Title sprites
+    function makeLabel(text){
+        const canvas = document.createElement('canvas'); const ctx = canvas.getContext('2d');
+        ctx.font = 'Bold 24px Arial'; const m = ctx.measureText(text);
+        canvas.width = m.width + 8; canvas.height = 40; ctx.font = 'Bold 24px Arial';
+        ctx.fillStyle = 'white'; ctx.fillText(text, 0, 28);
+        const tex = new THREE.CanvasTexture(canvas); const mat = new THREE.SpriteMaterial({ map: tex, transparent: true });
+        const sp = new THREE.Sprite(mat); sp.scale.set(4, 1.2, 1); return sp;
+    }
+    const labIn = makeLabel('Input'); labIn.position.set(-CUBE_SIZE*3.2, CUBE_SIZE*0.5, 0); inputGroup.add(labIn);
+    const labOut = makeLabel('Output (flipped)'); labOut.position.set(-CUBE_SIZE*3.2, -H*(CUBE_SIZE+GAP) - CUBE_SIZE*4.0, - (CUBE_SIZE*6)); outputGroup.add(labOut);
+
+    // Layout: separate layers in Z
+    inputGroup.position.set(0, 0, 0);
+    outputGroup.position.set(0, -H*(CUBE_SIZE+GAP) - CUBE_SIZE*4.0, 0);
+
+    // Interaction panel
+    const sideMenu = document.createElement('div');
+    Object.assign(sideMenu.style, { position:'absolute', top:'10px', right:'10px', width:'220px', padding:'10px', background:'rgba(0,0,0,0.7)', color:'#fff', borderRadius:'6px', zIndex:'2000' });
+    overlay.appendChild(sideMenu);
+    // Steps overlay toggle
+    const stepsBtn = document.createElement('button');
+    stepsBtn.textContent = 'Flip Steps 3D';
+    Object.assign(stepsBtn.style, { position:'absolute', top:'10px', left:'10px', zIndex:'2001' });
+    overlay.appendChild(stepsBtn);
+    let stepsCleanup = null;
+    stepsBtn.addEventListener('click', ()=>{
+        if (stepsCleanup){ try{ stepsCleanup(); }catch(e){} stepsCleanup=null; stepsBtn.textContent='Flip Steps 3D'; return; }
+        // choose length along flipped dimension if 2D; otherwise use total W
+        const n = is2D ? (dim===0? H: W) : W;
+        stepsCleanup = createFlip3D(overlay, { length: Math.min(128, n) });
+        stepsBtn.textContent = 'Close Steps';
+    });
+
+    async function getFlipValue(which, r, c){
+        const body = { uuid: op.uuid, which, x: c, y: r };
+        const res = await fetch(`${API_BASE}/api/getFlipValue`, { method:'POST', headers:{'Content-Type':'application/json'}, body: JSON.stringify(body)});
+        return await res.json();
+    }
+
+    const raycaster = new THREE.Raycaster();
+    const mouse = new THREE.Vector2();
+    let hoveredCube = null;
+    function clearHoverOutlines(){
+        // reset outlines on all cubes to avoid residual highlights
+        const groups = [inputGroup, outputGroup];
+        for (const g of groups){
+            if (!g) continue;
+            for (const child of g.children){
+                const ho = child && child.getObjectByName && child.getObjectByName('hoverOutline');
+                if (ho) ho.visible = false;
+            }
+        }
+    }
+    function syncHoverHighlight(){
+        // periodically enforce: only current hovered cube is highlighted
+        clearHoverOutlines();
+        if (hoveredCube){
+            const ho = hoveredCube.getObjectByName && hoveredCube.getObjectByName('hoverOutline');
+            if (ho) ho.visible = true;
+        }
+    }
+    // run every ~60ms (about 16 FPS) to avoid leaving residual highlights
+    const _hoverTimer = setInterval(syncHoverHighlight, 60);
+    function _updateMouseNDC(event) {
+        const rect = renderer.domElement.getBoundingClientRect();
+        const dpr = (window.devicePixelRatio || 1);
+        const px = (event.clientX - rect.left) * dpr;
+        const py = (event.clientY - rect.top ) * dpr;
+        const w = rect.width * dpr, h = rect.height * dpr;
+        mouse.x = (px / w) * 2 - 1; mouse.y = -(py / h) * 2 + 1;
+    }
+    function findTopLevel(obj){ let n=obj; while(n && !(n.userData && n.userData.tensorName)) n=n.parent; return n; }
+    function updatePanel(which, r, c, value){
+        if (!which){ sideMenu.innerHTML=''; return; }
+        sideMenu.innerHTML = `<h3>${which} Tensor</h3><p>Row: ${r+1}</p><p>Col: ${c+1}</p><p>Value: ${value !== undefined ? value : '...'}</p>`;
+    }
+
+    async function onMouseMove(event){
+        _updateMouseNDC(event); raycaster.setFromCamera(mouse, camera);
+        const objs = [...inputGroup.children, ...outputGroup.children];
+        const hits = raycaster.intersectObjects(objs, true);
+        // only update target; periodic timer will enforce single highlight
+        hoveredCube = null;
+        if (hits.length === 0){ updatePanel('',0,0,undefined); return; }
+        let cube = findTopLevel(hits[0].object); if (!cube) { updatePanel('',0,0,undefined); return; }
+        hoveredCube = cube;
+        const which = cube.userData.tensorName === 'Input' ? 'input' : 'output';
+        const r = cube.userData.tensor1 || 0; const c = cube.userData.tensor0 || 0;
+        try { const res = await getFlipValue(which, r, c); updatePanel(which, r, c, res.value); } catch(e){ updatePanel(which, r, c, undefined); }
+        // do not toggle here; timer will handle visibility to avoid residuals
+    }
+    overlay.addEventListener('mousemove', onMouseMove);
+    overlay.addEventListener('mouseleave', ()=>{ hoveredCube=null; updatePanel('',0,0,undefined); syncHoverHighlight(); });
+
+    // Camera framing
+    const box = new THREE.Box3().setFromObject(root); const center = box.getCenter(new THREE.Vector3()); const size = box.getSize(new THREE.Vector3());
+    const maxDim = Math.max(size.x, size.y, size.z); const fov = camera.fov * (Math.PI/180); let cameraZ = Math.abs(maxDim/2/Math.tan(fov/2)); cameraZ *= 1.6;
+    camera.position.set(center.x, center.y, center.z + cameraZ); camera.lookAt(center);
+
+    function animate(){ requestAnimationFrame(animate); controls.update(); renderer.render(scene, camera); }
+    animate();
+
+    const closeBtn = document.createElement('button'); closeBtn.textContent = 'Close'; Object.assign(closeBtn.style, { position:'absolute', top:'50px', left:'10px', zIndex:'2001' }); overlay.appendChild(closeBtn);
+    function cleanup(){ clearInterval(_hoverTimer); overlay.removeEventListener('mousemove', onMouseMove); overlay.removeEventListener('mouseleave', ()=>{}); if (stepsCleanup){ try{ stepsCleanup(); }catch(e){} stepsCleanup=null; } if (overlay && overlay.remove) overlay.remove(); }
+    closeBtn.addEventListener('click', cleanup);
+    try { window.current_op_uuid = op.uuid; } catch (e) {}
+    return cleanup;
+}
diff --git a/triton_viz/static/flip_3d.js b/triton_viz/static/flip_3d.js
new file mode 100644
index 00000000..c8351952
--- /dev/null
+++ b/triton_viz/static/flip_3d.js
@@ -0,0 +1,189 @@
+import * as THREE from 'https://esm.sh/three@0.155.0/build/three.module.js';
+import { OrbitControls } from 'https://esm.sh/three@0.155.0/examples/jsm/controls/OrbitControls.js';
+import { setupScene, setupGeometries, createCube, CUBE_SIZE, GAP } from './load_utils.js';
+
+// 3D Flip visualization with layered rows (initial + step_i view + step_i swap)
+// API: createFlip3D(container, { length, steps }) => cleanup()
+export function createFlip3D(containerElement, options) {
+    // overlay wrapper to guarantee visibility above existing content
+    const overlay = document.createElement('div');
+    Object.assign(overlay.style, {
+        position: 'absolute', left: '0', top: '0', right: '0', bottom: '0',
+        zIndex: 3000, pointerEvents: 'auto'
+    });
+    containerElement.appendChild(overlay);
+    const length = Math.max(2, (options && options.length) || 32);
+    const steps = (options && options.steps) || computeSteps(length);
+
+    const { scene, camera, renderer } = setupScene(overlay, 0x15151b);
+    const controls = new OrbitControls(camera, renderer.domElement);
+    controls.enableDamping = true; controls.dampingFactor = 0.06;
+
+    const { cubeGeometry, edgesGeometry, lineMaterial } = setupGeometries();
+
+    // Utilities
+    function hsv(h,s,v){
+        const c=v*s; const x=c*(1-Math.abs((h/60)%2-1)); const m=v-c; let r,g,b;
+        if(h<60){[r,g,b]=[c,x,0]} else if(h<120){[r,g,b]=[x,c,0]} else if(h<180){[r,g,b]=[0,c,x]}
+        else if(h<240){[r,g,b]=[0,x,c]} else if(h<300){[r,g,b]=[x,0,c]} else {[r,g,b]=[c,0,x]}
+        return new THREE.Color(r+m, g+m, b+m);
+    }
+    const hueAt = (val)=> 360 * (val/(length-1||1));
+
+    function createTextSprite(text){
+        const canvas = document.createElement('canvas');
+        const ctx = canvas.getContext('2d');
+        ctx.font = 'Bold 28px Arial';
+        const metrics = ctx.measureText(text);
+        canvas.width = metrics.width + 8; canvas.height = 40;
+        ctx.font = 'Bold 28px Arial';
+        ctx.fillStyle = 'white'; ctx.fillText(text, 0, 28);
+        const tex = new THREE.CanvasTexture(canvas);
+        const mat = new THREE.SpriteMaterial({ map: tex, transparent: true });
+        const sp = new THREE.Sprite(mat); sp.scale.set(4, 1.2, 1);
+        return sp;
+    }
+
+    // Build sequences with side-by-side panels per step
+    function computeSteps(n){ const arr=[]; let s=Math.floor(n/2); while(s>=1){ arr.push(s); s=Math.floor(s/2);} return arr; }
+    function swapBlocksArray(array, seg){
+        const res = array.slice(); const group = Math.max(1, 2*seg);
+        for(let g=0; g<res.length; g+=group){
+            const l = Math.min(seg, Math.max(0, res.length - g));
+            const r = Math.min(seg, Math.max(0, res.length - (g+seg)));
+            const len = Math.min(l, r);
+            for(let k=0;k<len;k++){
+                const i=g+k, j=g+seg+k; const t=res[i]; res[i]=res[j]; res[j]=t;
+            }
+        }
+        return res;
+    }
+
+    function layoutWidth(seg, n){
+        const groups = Math.ceil(n / (2*seg));
+        const perGroup = seg * (CUBE_SIZE+GAP);
+        const groupGap = CUBE_SIZE * 1.2;
+        return groups*perGroup + Math.max(0, groups-1)*groupGap;
+    }
+
+    function buildPanel(arrVals, seg, y, labelText, xOffset){
+        const panel = new THREE.Group();
+        const groupGap = CUBE_SIZE * 1.2;
+        const per = (CUBE_SIZE+GAP);
+        const groups = Math.ceil(arrVals.length / (2*seg));
+        for(let i=0;i<arrVals.length;i++){
+            const g = Math.floor(i / (2*seg));
+            const p = i % (2*seg);
+            const row = (p >= seg) ? 1 : 0;
+            const col = p % seg;
+            const color = hsv(hueAt(arrVals[i]), 1, 1);
+            const cube = createCube(color, 'Flip', col, row, g, cubeGeometry, edgesGeometry, lineMaterial);
+            const x = xOffset + g*(seg*per + groupGap) + col*per;
+            const yPos = y - row*per;
+            cube.position.set(x, yPos, 0);
+            panel.add(cube);
+        }
+        const label = createTextSprite(labelText);
+        label.position.set(xOffset - (CUBE_SIZE*3.2), y + CUBE_SIZE*0.5, 0);
+        panel.add(label);
+        return panel;
+    }
+
+    // Build all layers
+    const layers = [];
+    const root = new THREE.Group(); scene.add(root);
+
+    const rowGap = CUBE_SIZE * 4.2; // vertical distance per step row
+    const columnGap = CUBE_SIZE * 6.0; // horizontal gap between view and swap panels
+
+    let currentY = 0;
+    // initial 1D row (single panel)
+    const initPanel = new THREE.Group();
+    for (let i=0;i<length;i++){
+        const color = hsv(hueAt(i), 1, 1);
+        const cube = createCube(color, 'Flip', i, 0, 0, cubeGeometry, edgesGeometry, lineMaterial);
+        cube.position.set(i*(CUBE_SIZE+GAP), currentY, 0);
+        initPanel.add(cube);
+    }
+    const initLabel = createTextSprite('initial array');
+    initLabel.position.set(-CUBE_SIZE*3.2, currentY + CUBE_SIZE*0.5, 0);
+    initPanel.add(initLabel);
+    root.add(initPanel); layers.push(initPanel);
+
+    currentY -= rowGap;
+    let arr = Array.from({length}, (_,i)=>i);
+    for (let s=0; s<steps.length; s++){
+        const seg = steps[s];
+        const lw = layoutWidth(seg, length);
+        const viewX = 0;
+        const swapX = viewX + lw + columnGap;
+        const viewPanel = buildPanel(arr, seg, currentY, `step ${s} view`, viewX);
+        const swapped = swapBlocksArray(arr, seg);
+        const swapPanel = buildPanel(swapped, seg, currentY, `step ${s} swap`, swapX);
+        root.add(viewPanel); root.add(swapPanel);
+        layers.push(viewPanel); layers.push(swapPanel);
+        arr = swapped.slice();
+        currentY -= rowGap;
+    }
+
+    // final reshape (1D row)
+    const finalPanel = new THREE.Group();
+    for (let i=0;i<arr.length;i++){
+        const color = hsv(hueAt(arr[i]), 1, 1);
+        const cube = createCube(color, 'Flip', i, 0, 0, cubeGeometry, edgesGeometry, lineMaterial);
+        cube.position.set(i*(CUBE_SIZE+GAP), currentY, 0);
+        finalPanel.add(cube);
+    }
+    const finalLabel = createTextSprite('final reshape');
+    finalLabel.position.set(-CUBE_SIZE*3.2, currentY + CUBE_SIZE*0.5, 0);
+    finalPanel.add(finalLabel);
+    root.add(finalPanel); layers.push(finalPanel);
+
+    // Set initial visibilities
+    layers.forEach((g, i)=> g.visible = (i===0));
+
+    // Center camera to see all layers
+    const totalW = Math.max(layoutWidth(steps[0]||Math.floor(length/2), length) + columnGap + layoutWidth(steps[0]||Math.floor(length/2), length), length*(CUBE_SIZE+GAP));
+    const totalH = (layers.length) * rowGap;
+    camera.position.set(totalW*0.6, -totalH*0.35, Math.max(6, totalW*0.9));
+    camera.lookAt(new THREE.Vector3(totalW*0.5, -totalH*0.4, 0));
+
+    // UI controls
+    const ui = document.createElement('div');
+    Object.assign(ui.style, { position: 'absolute', top: '10px', left: '10px', display: 'flex', gap: '8px', zIndex: 2000 });
+    const playBtn = document.createElement('button'); playBtn.textContent = 'Play';
+    const stepLabel = document.createElement('span'); stepLabel.textContent = 'Layer: 1/' + layers.length; stepLabel.style.color = '#fff';
+    const speedSel = document.createElement('select'); ['0.5x','1x','2x','4x'].forEach(s=>{ const o=document.createElement('option'); o.value=s; o.textContent=s; if(s==='1x') o.selected=true; speedSel.appendChild(o); });
+    ui.appendChild(playBtn); ui.appendChild(stepLabel); ui.appendChild(speedSel); overlay.appendChild(ui);
+    const closeBtn = document.createElement('button'); closeBtn.textContent = 'Close';
+    ui.appendChild(closeBtn);
+
+    // Animation loop: reveal rows sequentially
+    let playing=false; let lastTs=0; let idxVisible=0; let raf=0;
+    function rate(){ return speedSel.value==='0.5x'?800: speedSel.value==='1x'?400: speedSel.value==='2x'?200:100; }
+    function loop(){
+        controls.update();
+        if (playing && (lastTs===0 || performance.now()-lastTs>=rate())){
+            if (idxVisible < layers.length-1){
+                idxVisible += 1; layers[idxVisible].visible = true; stepLabel.textContent = `Layer: ${idxVisible+1}/${layers.length}`; lastTs = performance.now();
+            } else { playing=false; playBtn.textContent='Replay'; }
+        }
+        renderer.render(scene, camera);
+        raf = requestAnimationFrame(loop);
+    }
+    raf = requestAnimationFrame(loop);
+
+    playBtn.addEventListener('click', ()=>{
+        if (!playing && idxVisible===layers.length-1){
+            // reset visibility
+            layers.forEach((r, i)=> r.visible = (i===0)); idxVisible = 0; stepLabel.textContent = `Layer: 1/${layers.length}`;
+        }
+        playing = !playing; playBtn.textContent = playing? 'Pause' : (idxVisible===layers.length-1? 'Replay' : 'Play');
+        if (playing) lastTs = 0;
+    });
+
+    closeBtn.addEventListener('click', ()=> cleanup());
+
+    function cleanup(){ cancelAnimationFrame(raf); controls.dispose(); renderer.dispose?.(); if (overlay && overlay.remove) overlay.remove(); }
+    return cleanup;
+}
diff --git a/triton_viz/static/flip_demo.js b/triton_viz/static/flip_demo.js
new file mode 100644
index 00000000..34ecc5d9
--- /dev/null
+++ b/triton_viz/static/flip_demo.js
@@ -0,0 +1,178 @@
+// Lightweight 2D canvas demo to visualize stepwise flip via view/reshape + swap
+// API: createFlipDemo(containerElement, options) -> cleanup()
+// options: { length: number, steps?: number[] } where steps are segment sizes halving each step
+
+export function createFlipDemo(containerElement, options) {
+    const length = Math.max(1, (options && options.length) || 16);
+    const steps = (options && options.steps) || computeSteps(length);
+
+    const wrapper = document.createElement('div');
+    Object.assign(wrapper.style, {
+        position: 'absolute', left: '0', top: '0', right: '0', bottom: '0',
+        display: 'flex', flexDirection: 'column', gap: '10px', padding: '10px'
+    });
+
+    const toolbar = document.createElement('div');
+    Object.assign(toolbar.style, { display: 'flex', gap: '8px', alignItems: 'center' });
+    const playBtn = document.createElement('button'); playBtn.textContent = 'Play';
+    const stepLabel = document.createElement('span'); stepLabel.textContent = 'Step: 0';
+    const speedSel = document.createElement('select');
+    ;['0.5x','1x','2x','4x'].forEach(s=>{ const o=document.createElement('option'); o.value=s; o.textContent=s; if(s==='1x') o.selected=true; speedSel.appendChild(o); });
+    toolbar.appendChild(playBtn); toolbar.appendChild(stepLabel); toolbar.appendChild(speedSel);
+
+    const canvas = document.createElement('canvas');
+    canvas.width = Math.min(1200, containerElement.clientWidth - 20);
+    canvas.height = Math.min(500, containerElement.clientHeight - 20);
+    const ctx = canvas.getContext('2d');
+
+    wrapper.appendChild(toolbar);
+    wrapper.appendChild(canvas);
+    containerElement.appendChild(wrapper);
+
+    // data: 0..length-1
+    let data = new Array(length).fill(0).map((_,i)=>i);
+    let currentStep = 0;
+    let playing = false;
+    let rafId = 0;
+
+    function computeSteps(n){
+        // e.g. for 16 -> [8,4,2,1]
+        const arr=[]; let s = Math.floor(n/2);
+        while(s>=1){ arr.push(s); s=Math.floor(s/2);} return arr;
+    }
+
+    function lerp(a,b,t){ return a+(b-a)*t; }
+
+    function hsv(h,s,v){
+        let c=v*s; let x=c*(1-Math.abs((h/60)%2-1)); let m=v-c; let r,g,b;
+        if(h<60){[r,g,b]=[c,x,0]} else if(h<120){[r,g,b]=[x,c,0]} else if(h<180){[r,g,b]=[0,c,x]}
+        else if(h<240){[r,g,b]=[0,x,c]} else if(h<300){[r,g,b]=[x,0,c]} else {[r,g,b]=[c,0,x]}
+        return `rgb(${Math.round((r+m)*255)},${Math.round((g+m)*255)},${Math.round((b+m)*255)})`;
+    }
+
+    function drawArray(arr, y, boxW, boxH, highlightPairs){
+        for(let i=0;i<arr.length;i++){
+            const x = 20 + i*boxW;
+            const hue = lerp(0, 350, i/(arr.length-1||1));
+            ctx.fillStyle = hsv(hue, 1, 1);
+            ctx.strokeStyle = '#0a0a1a';
+            ctx.lineWidth = 2;
+            ctx.fillRect(x, y, boxW-2, boxH);
+            ctx.strokeRect(x, y, boxW-2, boxH);
+            ctx.fillStyle = '#fff';
+            ctx.font = '14px monospace';
+            ctx.textAlign = 'center'; ctx.textBaseline = 'middle';
+            ctx.fillText(String(arr[i]), x + (boxW-2)/2, y + boxH/2);
+        }
+        if (highlightPairs && highlightPairs.length){
+            ctx.strokeStyle = 'rgba(255,255,255,0.8)';
+            highlightPairs.forEach(([i,j])=>{
+                const x1 = 20 + i*boxW + (boxW-2)/2;
+                const x2 = 20 + j*boxW + (boxW-2)/2;
+                const yy = y + boxH + 6;
+                ctx.beginPath(); ctx.moveTo(x1, yy); ctx.lineTo(x2, yy); ctx.stroke();
+            });
+        }
+    }
+
+    function calcPairs(seg){
+        // pair indices across two adjacent blocks of size `seg` inside each group of size `2*seg`
+        const pairs=[]; const n=length; const group = Math.max(1, 2*seg);
+        for(let g=0; g<n; g+=group){
+            const leftStart = g;
+            const rightStart = g + seg;
+            const leftLen = Math.min(seg, Math.max(0, n - leftStart));
+            const rightLen = Math.min(seg, Math.max(0, n - rightStart));
+            const pairLen = Math.min(leftLen, rightLen);
+            for(let k=0;k<pairLen;k++){
+                pairs.push([leftStart + k, rightStart + k]);
+            }
+        }
+        return pairs;
+    }
+
+    function performSwap(arr, seg){
+        // swap neighboring blocks of length `seg` within each group of length `2*seg`
+        const res = arr.slice();
+        const n = res.length; const group = Math.max(1, 2*seg);
+        for(let g=0; g<n; g+=group){
+            const leftStart = g;
+            const rightStart = g + seg;
+            const leftLen = Math.min(seg, Math.max(0, n - leftStart));
+            const rightLen = Math.min(seg, Math.max(0, n - rightStart));
+            const len = Math.min(leftLen, rightLen);
+            for(let k=0;k<len;k++){
+                const i = leftStart + k, j = rightStart + k;
+                const tmp=res[i]; res[i]=res[j]; res[j]=tmp;
+            }
+        }
+        return res;
+    }
+
+    function drawFrame(t){
+        ctx.clearRect(0,0,canvas.width,canvas.height);
+        ctx.fillStyle='#1e1e28'; ctx.fillRect(0,0,canvas.width,canvas.height);
+        const boxW = Math.max(18, Math.min(42, (canvas.width-40)/length));
+        const boxH = 26;
+
+        // initial array
+        ctx.fillStyle = '#fff'; ctx.font = '16px sans-serif'; ctx.fillText('initial array', 20, 24);
+        drawArray(new Array(length).fill(0).map((_,i)=>i), 40, boxW, boxH);
+
+        let y = 100;
+        for(let s=0;s<steps.length;s++){
+            const seg = steps[s];
+            ctx.fillStyle = '#fff'; ctx.font='14px sans-serif';
+            ctx.fillText(`step ${s} view`, 20, y-6);
+            drawArray(data, y, boxW, boxH);
+            const pairs = calcPairs(seg);
+            ctx.fillText(`step ${s} swap`, 20 + Math.min(canvas.width-20, (length*boxW)/2), y-6);
+            // Draw swapped preview to the right
+            const swapped = performSwap(data, seg);
+            const offsetX = Math.min(canvas.width-20, (length*boxW)/2);
+            ctx.save(); ctx.translate(offsetX, 0);
+            drawArray(swapped, y, boxW, boxH, pairs);
+            ctx.restore();
+            y += boxH + 60;
+            if (s === currentStep) break; // show up to current step during animation
+        }
+
+        if (currentStep === steps.length) {
+            ctx.fillStyle = '#fff'; ctx.font='16px sans-serif'; ctx.fillText('final reshape', 20, y-6);
+            drawArray(data, y, boxW, boxH);
+        }
+    }
+
+    function tick(){
+        drawFrame();
+        if (!playing) return;
+        const rate = speedSel.value==='0.5x'?800: speedSel.value==='1x'?400: speedSel.value==='2x'?200:100; // ms per step
+        if (_lastStepTime === 0 || Date.now() - _lastStepTime >= rate){
+            if (currentStep < steps.length){
+                data = performSwap(data, steps[currentStep]);
+                currentStep += 1;
+                stepLabel.textContent = `Step: ${currentStep}`;
+                _lastStepTime = Date.now();
+            } else {
+                playing = false; playBtn.textContent = 'Replay';
+            }
+        }
+        rafId = requestAnimationFrame(tick);
+    }
+    let _lastStepTime = 0;
+
+    playBtn.addEventListener('click', ()=>{
+        if (!playing && currentStep===steps.length){
+            // reset
+            data = new Array(length).fill(0).map((_,i)=>i);
+            currentStep = 0; stepLabel.textContent='Step: 0';
+        }
+        playing = !playing; playBtn.textContent = playing ? 'Pause' : (currentStep===steps.length ? 'Replay' : 'Play');
+        if (playing){ _lastStepTime = 0; cancelAnimationFrame(rafId); rafId = requestAnimationFrame(tick);}
+    });
+
+    // initial paint
+    drawFrame(0);
+
+    return function cleanup(){ cancelAnimationFrame(rafId); if (wrapper && wrapper.remove) wrapper.remove(); };
+}
diff --git a/triton_viz/static/gridblock.js b/triton_viz/static/gridblock.js
index 7ee9762f..89081726 100644
--- a/triton_viz/static/gridblock.js
+++ b/triton_viz/static/gridblock.js
@@ -1,4 +1,5 @@
 import { createMatMulVisualization } from './matmul.js';
+import { createFlipVisualization } from './flip.js';
 import { createLoadVisualization } from './load.js';
 import { createStoreVisualization } from './store.js';
 
@@ -74,6 +75,9 @@ export class GridBlock {
 
         const closeButton = this.createCloseButton();
         this.visualizationContainer.appendChild(closeButton);
+        // Add an explicit Back button (top-left) to return to main canvas view
+        const backButton = this.createBackButton();
+        this.visualizationContainer.appendChild(backButton);
 
         // Ensure buttons panel sits above the canvas and accepts clicks
         const buttonsPanel = this.visualizationContainer.querySelector('div');
@@ -121,7 +125,8 @@ export class GridBlock {
                 header.style.opacity = '0.9';
                 wrapper.appendChild(header);
                 try {
-                    const res = await fetch('/api/op_code', { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ uuid, frame_idx: frameIdx, context }) });
+                    const API_BASE = window.__TRITON_VIZ_API__ || '';
+                    const res = await fetch(`${API_BASE}/api/op_code`, { method: 'POST', headers: { 'Content-Type': 'application/json' }, body: JSON.stringify({ uuid, frame_idx: frameIdx, context }) });
                     const data = await res.json();
                     const meta = document.createElement('div');
                     meta.style.marginBottom = '4px';
@@ -273,6 +278,9 @@ export class GridBlock {
             case 'Store':
                 this.visualizationCleanupFunction = createStoreVisualization(this.contentArea, op);
                 break;
+            case 'Flip':
+                this.visualizationCleanupFunction = createFlipVisualization(this.contentArea, op);
+                break;
             default:
                 const unsupportedMsg = document.createElement('p');
                 unsupportedMsg.textContent = `Visualization not supported for ${op.type} operation`;
@@ -295,6 +303,26 @@ export class GridBlock {
         return closeButton;
     }
 
+    createBackButton() {
+        const backButton = document.createElement('button');
+        backButton.textContent = 'Back';
+        Object.assign(backButton.style, {
+            position: 'fixed',
+            left: '10px',
+            bottom: '10px',
+            zIndex: '2002',
+            background: 'rgba(0,0,0,0.65)',
+            color: '#fff',
+            border: '1px solid #666',
+            padding: '6px 10px',
+            borderRadius: '6px',
+            cursor: 'pointer',
+            boxShadow: '0 2px 8px rgba(0,0,0,0.5)'
+        });
+        backButton.addEventListener('click', () => this.hideDetailedView());
+        return backButton;
+    }
+
     hideDetailedView() {
         if (!this.isDetailedViewVisible) return;
 
diff --git a/triton_viz/static/load.js b/triton_viz/static/load.js
index 25274c3f..9e08f10d 100644
--- a/triton_viz/static/load.js
+++ b/triton_viz/static/load.js
@@ -14,7 +14,8 @@ import {
 export function createLoadVisualization(containerElement, op) {
 
         console.log(op.uuid);
-        fetch('/api/setop', {
+        const API_BASE = window.__TRITON_VIZ_API__ || '';
+        fetch(`${API_BASE}/api/setop`, {
             method: 'POST',
             headers: {
                 'Content-Type': 'application/json',
@@ -351,7 +352,7 @@ export function createLoadVisualization(containerElement, op) {
             wrapper.appendChild(header);
 
             try {
-                const res = await fetch('/api/op_code', {
+                const res = await fetch(`${API_BASE}/api/op_code`, {
                     method: 'POST',
                     headers: { 'Content-Type': 'application/json' },
                     body: JSON.stringify({ uuid: op.uuid, frame_idx: frameIdx, context })
@@ -497,7 +498,7 @@ export function createLoadVisualization(containerElement, op) {
 
         async function getElementValue(tensorName, x, y, z) {
             let uuid = op.uuid;
-            const response = await fetch('/api/getLoadValue', {
+            const response = await fetch(`${API_BASE}/api/getLoadValue`, {
                 method: 'POST',
                 headers: {
                     'Content-Type': 'application/json',
@@ -509,7 +510,7 @@ export function createLoadVisualization(containerElement, op) {
 
         async function fetchGlobalTensor() {
             try {
-                const res = await fetch('/api/getLoadTensor', {
+                const res = await fetch(`${API_BASE}/api/getLoadTensor`, {
                     method: 'POST',
                     headers: { 'Content-Type': 'application/json' },
                     body: JSON.stringify({ uuid: op.uuid })
diff --git a/triton_viz/static/matmul.js b/triton_viz/static/matmul.js
index 3ce14757..4db64f8e 100644
--- a/triton_viz/static/matmul.js
+++ b/triton_viz/static/matmul.js
@@ -2,9 +2,10 @@ import * as THREE from 'https://esm.sh/three@0.155.0/build/three.module.js';
 import { OrbitControls } from 'https://esm.sh/three@0.155.0/examples/jsm/controls/OrbitControls.js';
 
 export function createMatMulVisualization(containerElement, op) {
+    const API_BASE = window.__TRITON_VIZ_API__ || '';
     const { input_shape, other_shape, output_shape } = op;
     console.log(op.uuid)
-    fetch('/api/setop', {
+    fetch(`${API_BASE}/api/setop`, {
         method: 'POST',
         headers: {
             'Content-Type': 'application/json',
@@ -14,9 +15,7 @@ export function createMatMulVisualization(containerElement, op) {
     .then(response => response.json())
     .then(data => console.log('Set current op:', data))
     .catch((error) => console.error('Error:', error));
-    let currentStep = 0;
-    const totalSteps = input_shape[1];
-    let frame = 0;
+    let currentStep = 0; // kept for compatibility with getValue; not used for animation
 
 
     const sideMenu = document.createElement('div');
@@ -30,6 +29,8 @@ export function createMatMulVisualization(containerElement, op) {
     sideMenu.style.fontFamily = 'Arial, sans-serif';
     sideMenu.style.fontSize = '14px';
     sideMenu.style.borderRadius = '5px';
+    sideMenu.style.zIndex = '2000';
+    sideMenu.style.pointerEvents = 'auto';
     containerElement.appendChild(sideMenu);
     let hoveredCube = null;
 
@@ -38,7 +39,7 @@ export function createMatMulVisualization(containerElement, op) {
 
     async function getElementValue( matrixName, row, col) {
         let uuid = op.uuid;
-        const response = await fetch('/api/getValue', {
+        const response = await fetch(`${API_BASE}/api/getValue`, {
             method: 'POST',
             headers: {
                 'Content-Type': 'application/json',
@@ -49,7 +50,7 @@ export function createMatMulVisualization(containerElement, op) {
     }
 
 
-    function updateSideMenu(matrix, x, y) {
+    function updateSideMenu(matrix, x, y, vectors) {
         if (!matrix) {
             sideMenu.innerHTML = '';
             return;
@@ -71,11 +72,25 @@ export function createMatMulVisualization(containerElement, op) {
             return;
         }
         console.log(matrixName, "x:", (x + 1), "y:", (y + 1));
+        let extra = '';
+        if (matrixName === 'C' && vectors && !vectors.error) {
+            const aRow = vectors.a_row || [];
+            const bCol = vectors.b_col || [];
+            const k = vectors.k || 0;
+            extra = `
+                <hr/>
+                <div><b>From A row</b>: [${aRow.slice(0,8).join(', ')}${aRow.length>8?' …':''}]</div>
+                <div><b>From B col</b>: [${bCol.slice(0,8).join(', ')}${bCol.length>8?' …':''}]</div>
+                <div><b>k</b>: ${k}</div>
+            `;
+        }
+
         sideMenu.innerHTML = `
             <h3 style="margin-top: 0;">Matrix ${matrixName}</h3>
             <p>Row: ${y + 1}</p>
             <p>Column: ${x + 1}</p>
             <p>Dimensions: ${dims[0]} x ${dims[1]}</p>
+            ${extra}
         `;
     }
 
@@ -135,10 +150,37 @@ export function createMatMulVisualization(containerElement, op) {
                     `Value: ${res.value}`
                 );
 
-                updateSideMenu(hoveredCube.matrixName, hoveredCube.matrixRow, hoveredCube.matrixCol);
+                // If hovering C, fetch vectors and update panel
+                let vectors = null;
+                if (hoveredCube.matrixName === 'C') {
+                    try {
+                        const resp = await fetch(`${API_BASE}/api/getMatmulVectors`, {
+                            method: 'POST',
+                            headers: { 'Content-Type': 'application/json' },
+                            body: JSON.stringify({ uuid: op.uuid, row: hoveredCube.matrixRow, col: hoveredCube.matrixCol })
+                        });
+                        vectors = await resp.json();
+                    } catch (e) { vectors = { error: String(e) }; }
+
+                    // highlight A[row,:], B[:,col], C[row,col]
+                    resetColors();
+                    const row = hoveredCube.matrixRow;
+                    const col = hoveredCube.matrixCol;
+                    const highlightA = Array.from({ length: input_shape[1] }, (_, i) => [row, i]);
+                    const highlightB = Array.from({ length: other_shape[0] }, (_, i) => [i, col]);
+                    const highlightC = [[row, col]];
+                    highlightCubes(matrixA, highlightA, COLOR_HIGHLIGHT);
+                    highlightCubes(matrixB, highlightB, COLOR_HIGHLIGHT);
+                    highlightCubes(matrixC, highlightC, COLOR_FILLED);
+                } else {
+                    // hovering A/B or others -> just reset
+                    resetColors();
+                }
+                updateSideMenu(hoveredCube.matrixName, hoveredCube.matrixRow, hoveredCube.matrixCol, vectors);
             }
         } else {
             updateSideMenu(null);
+            resetColors();
         }
     }
 
@@ -242,9 +284,7 @@ export function createMatMulVisualization(containerElement, op) {
     controls.target.copy(center);
     controls.update();
 
-    let isPaused = false;
-
-    const totalFrames = input_shape[0] * other_shape[1];
+    // remove auto animation; highlight is driven by mouse hover only
 
     function highlightCubes(matrix, indices, highlightColor) {
         indices.forEach(([i, j]) => {
@@ -267,25 +307,6 @@ export function createMatMulVisualization(containerElement, op) {
     function animate() {
         requestAnimationFrame(animate);
         controls.update();
-
-        if (!isPaused && frame < totalFrames) {
-            resetColors();
-
-            const row = Math.floor(frame / other_shape[1]);
-            const col = frame % other_shape[1];
-            currentStep = frame % totalSteps + 1;
-
-            const highlightA = Array.from({ length: input_shape[1] }, (_, i) => [row, i]);
-            const highlightB = Array.from({ length: other_shape[0] }, (_, i) => [i, col]);
-            const highlightC = [[row, col]];
-
-            highlightCubes(matrixA, highlightA, COLOR_HIGHLIGHT);
-            highlightCubes(matrixB, highlightB, COLOR_HIGHLIGHT);
-            highlightCubes(matrixC, highlightC, COLOR_FILLED);
-
-            frame++;
-        }
-
         renderer.render(scene, camera);
     }
 
@@ -301,22 +322,10 @@ export function createMatMulVisualization(containerElement, op) {
     controlPanel.style.left = '10px';
     controlPanel.style.display = 'flex';
     controlPanel.style.gap = '10px';
+    controlPanel.style.zIndex = '2000';
+    controlPanel.style.pointerEvents = 'auto';
 
-    const playPauseButton = document.createElement('button');
-    playPauseButton.textContent = 'Play/Pause';
-    playPauseButton.addEventListener('click', () => {
-        isPaused = !isPaused;
-    });
-
-    const resetButton = document.createElement('button');
-    resetButton.textContent = 'Reset';
-    resetButton.addEventListener('click', () => {
-        frame = 0;
-        resetColors();
-    });
-
-    controlPanel.appendChild(playPauseButton);
-    controlPanel.appendChild(resetButton);
+    // Removed animation controls; keep panel for color-by-value toggle only
 
     // Color by Value (for C matrix) toggle + legend (mono blue)
     const colorToggle = document.createElement('button');
@@ -328,7 +337,7 @@ export function createMatMulVisualization(containerElement, op) {
     function createLegend(min,max){
         destroyLegend();
         const w = document.createElement('div');
-        Object.assign(w.style,{position:'absolute', left:'10px', bottom:'60px', background:'rgba(0,0,0,0.6)', color:'#fff', padding:'6px 8px', borderRadius:'6px'});
+        Object.assign(w.style,{position:'absolute', left:'10px', bottom:'60px', background:'rgba(0,0,0,0.6)', color:'#fff', padding:'6px 8px', borderRadius:'6px', zIndex:'2000', pointerEvents:'auto'});
         const c = document.createElement('canvas'); c.width=220; c.height=10; const ctx=c.getContext('2d');
         for(let x=0;x<c.width;x++){ const t=x/(c.width-1); const r=t,g=0.2,b=1-t; ctx.fillStyle=`rgb(${Math.round(r*255)},${Math.round(g*255)},${Math.round(b*255)})`; ctx.fillRect(x,0,1,c.height);}
         const lab=document.createElement('div'); lab.style.display='flex'; lab.style.justifyContent='space-between'; lab.style.marginTop='2px'; lab.innerHTML=`<span>${min.toFixed(3)}</span><span>${max.toFixed(3)}</span>`;
@@ -336,7 +345,7 @@ export function createMatMulVisualization(containerElement, op) {
         w.appendChild(ttl); w.appendChild(c); w.appendChild(lab); containerElement.appendChild(w); legendEl=w;
     }
     async function fetchCValues(){
-        try{ const res=await fetch('/api/getMatmulC',{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({uuid: op.uuid})});
+        try{ const res=await fetch(`${API_BASE}/api/getMatmulC`,{method:'POST',headers:{'Content-Type':'application/json'},body:JSON.stringify({uuid: op.uuid})});
             return await res.json(); }catch(e){ return {error:String(e)} }
     }
     function applyColorCWithData(data){ if(!colorOn||!data||data.error) return; try{
diff --git a/triton_viz/static/store-utils.js b/triton_viz/static/store-utils.js
index d87d78bc..12b9df11 100644
--- a/triton_viz/static/store-utils.js
+++ b/triton_viz/static/store-utils.js
@@ -3,7 +3,8 @@ import * as THREE from 'https://cdn.jsdelivr.net/npm/three@0.155.0/build/three.m
 export function createMatMulVisualization(containerElement, op) {
     const { input_shape, other_shape, output_shape } = op;
     console.log(op.uuid)
-    fetch('/api/setop', {
+    const API_BASE = window.__TRITON_VIZ_API__ || '';
+    fetch(`${API_BASE}/api/setop`, {
         method: 'POST',
         headers: {
             'Content-Type': 'application/json',
@@ -37,7 +38,7 @@ export function createMatMulVisualization(containerElement, op) {
 
     async function getElementValue( matrixName, row, col) {
         let uuid = op.uuid;
-        const response = await fetch('/api/getValue', {
+        const response = await fetch(`${API_BASE}/api/getValue`, {
             method: 'POST',
             headers: {
                 'Content-Type': 'application/json',
diff --git a/triton_viz/static/store.js b/triton_viz/static/store.js
index 9868c16a..cfca8838 100644
--- a/triton_viz/static/store.js
+++ b/triton_viz/static/store.js
@@ -10,11 +10,14 @@ import {
     setupEventListeners,
     cameraControls
 } from './load_utils.js';
+import { createFlipDemo } from './flip_demo.js';
+import { createFlip3D } from './flip_3d.js';
 
 export function createStoreVisualization(containerElement, op) {
 
         console.log(op.uuid);
-        fetch('/api/setop', {
+        const API_BASE = window.__TRITON_VIZ_API__ || '';
+        fetch(`${API_BASE}/api/setop`, {
             method: 'POST',
             headers: {
                 'Content-Type': 'application/json',
@@ -45,6 +48,7 @@ export function createStoreVisualization(containerElement, op) {
         containerElement.appendChild(controlBar);
         let dragModeOn = false;
         let hoveredCube = null;
+        let flipCleanup = null;
 
         const COLOR_GLOBAL = new THREE.Color(0.2, 0.2, 0.2);    // Dark Gray
         const COLOR_SLICE = new THREE.Color(0.0, 0.7, 1.0);     // Cyan (starting color for global slice)
@@ -95,6 +99,7 @@ export function createStoreVisualization(containerElement, op) {
             dragToggle.textContent = `Drag Cubes: ${dragModeOn ? 'ON' : 'OFF'}`;
             orbitControls.enabled = !dragModeOn;
         });
+        // Removed Flip demo button from Store view; Flip visualization is available under Flip op.
         animate();
 
         function _updateMouseNDC(event) {
@@ -229,7 +234,7 @@ export function createStoreVisualization(containerElement, op) {
 
         async function getElementValue(tensorName, x, y, z) {
             let uuid = op.uuid;
-            const response = await fetch('/api/getLoadValue', {
+            const response = await fetch(`${API_BASE}/api/getLoadValue`, {
                 method: 'POST',
                 headers: {
                     'Content-Type': 'application/json',
diff --git a/triton_viz/static/visualization.js b/triton_viz/static/visualization.js
index 15738ea5..bbfe8deb 100644
--- a/triton_viz/static/visualization.js
+++ b/triton_viz/static/visualization.js
@@ -379,7 +379,8 @@ function draw() {
 
 async function fetchData() {
     try {
-        const response = await fetch('/api/data');
+        const API_BASE = window.__TRITON_VIZ_API__ || '';
+        const response = await fetch(`${API_BASE}/api/data`);
         globalData = await response.json();
         console.log(globalData);
 
diff --git a/triton_viz/visualizer/draw.py b/triton_viz/visualizer/draw.py
index 22dd6a1b..6ea66c58 100644
--- a/triton_viz/visualizer/draw.py
+++ b/triton_viz/visualizer/draw.py
@@ -5,6 +5,7 @@
     Dot,
     Load,
     Store,
+    Flip,
 )
 from triton_viz.clients.sanitizer.data import OutOfBoundsRecordBruteForce
 import numpy as np
@@ -190,6 +191,40 @@ def prepare_visualization_data(program_records, tensor_table):
                 # prepare C values after kernel (if available from intermediate or recompute best-effort)
             }
 
+        elif isinstance(record, Flip):
+            visualization_data.append(
+                {
+                    "type": "Flip",
+                    "input_shape": record.input_shape,
+                    "output_shape": record.output_shape,
+                    "dim": int(getattr(record, "dim", 0)),
+                    "uuid": record_uuid,
+                }
+            )
+
+            raw_tensor_data[record_uuid] = {
+                "tracebacks": [
+                    {
+                        "filename": f.filename,
+                        "lineno": f.lineno,
+                        "line": f.line,
+                        "name": f.name,
+                    }
+                    for f in getattr(record, "call_path", [])
+                ],
+                # best-effort payload for potential future value viz
+                "input_shape": list(record.input_shape),
+                "output_shape": list(record.output_shape),
+                "dim": int(getattr(record, "dim", 0)),
+                # optionally include data for hover value queries
+                "input_data": None
+                if getattr(record, "input_data", None) is None
+                else torch.tensor(record.input_data),
+                "output_data": None
+                if getattr(record, "output_data", None) is None
+                else torch.tensor(record.output_data),
+            }
+
         elif isinstance(record, Load):
             global_tensor, slice_tensor = tensor_table[record.ptr]
             print(global_tensor)
diff --git a/triton_viz/visualizer/interface.py b/triton_viz/visualizer/interface.py
index 95b4f0de..c287b5ef 100644
--- a/triton_viz/visualizer/interface.py
+++ b/triton_viz/visualizer/interface.py
@@ -203,10 +203,18 @@ def get_op_code():
 
     def _score(tb: dict) -> int:
         fn = tb.get("filename") or ""
-        line = tb.get("line") or ""
+        line = (tb.get("line") or "").strip()
         name = tb.get("name") or ""
         p = os.path.realpath(fn)
         score = 0
+
+        # 强优先：明确的 Triton 语言内核操作符
+        if "tl.load" in line or "tl.store" in line:
+            score += 100
+        elif "tl." in line:
+            score += 50
+
+        # 位置相关：项目代码优先，三方/框架降权
         if p.startswith(cwd):
             score += 5
         if any(
@@ -214,12 +222,15 @@ def _score(tb: dict) -> int:
             for s in ["site-packages", "triton_viz/", "triton/", "runpy.py", "IPython"]
         ):
             score -= 10
+
+        # 语义相关：函数名看起来像 kernel 的加分
         if name.endswith("_kernel") or "kernel" in name:
             score += 3
-        if "tl." in line:
-            score += 2
+
+        # 体验相关：examples 目录小幅加分
         if "examples" in p:
             score += 1
+
         return score
 
     # Prefer frames from tail (closest to current) and highest score
@@ -287,6 +298,40 @@ def get_matmul_c():
         return jsonify({"error": f"MatMul compute failed: {e}"}), 200
 
 
+@app.route("/api/getMatmulVectors", methods=["POST"])
+def get_matmul_vectors():
+    """Return A[row, :] and B[:, col] for a given Dot op.
+
+    Request JSON: { uuid, row, col }
+    Response JSON: {
+        "row": int,
+        "col": int,
+        "a_row": [float, ...],
+        "b_col": [float, ...],
+        "k": int
+    }
+    """
+    global raw_tensor_data
+    data = request.json or {}
+    uuid = data.get("uuid")
+    row = int(data.get("row", 0))
+    col = int(data.get("col", 0))
+    if not uuid or uuid not in raw_tensor_data:
+        return jsonify({"error": "Operation not found"}), 404
+    op = raw_tensor_data[uuid]
+    a = op.get("input_data")
+    b = op.get("other_data")
+    if a is None or b is None:
+        return jsonify({"error": "MatMul tensors not available"}), 200
+    try:
+        a_row = a[row, :].cpu().numpy().tolist()
+        b_col = b[:, col].cpu().numpy().tolist()
+        k = len(a_row)
+        return jsonify({"row": row, "col": col, "a_row": a_row, "b_col": b_col, "k": k})
+    except Exception as e:
+        return jsonify({"error": f"MatMul vectors failed: {e}"}), 200
+
+
 @app.route("/api/getValue", methods=["POST"])
 def get_value():
     global raw_tensor_data, precomputed_c_values, current_fullscreen_op
@@ -364,6 +409,43 @@ def get_load_value():
         return jsonify({"error": "Global tensor data not found"}), 200
 
 
+@app.route("/api/getFlipValue", methods=["POST"])
+def get_flip_value():
+    """Return a value from input or output of a Flip op by linear index or 2D coords.
+
+    Request JSON: { uuid, which: "input"|"output", x, y }
+    - If tensor is 1D, use x only; if 2D, use x (col), y (row)
+    """
+    global raw_tensor_data
+    data = request.json or {}
+    uuid = data.get("uuid")
+    which = (data.get("which") or "input").lower()
+    x = data.get("x")
+    y = data.get("y")
+    if not uuid or uuid not in raw_tensor_data:
+        return jsonify({"error": "Operation not found"}), 404
+    op = raw_tensor_data[uuid]
+    t = op.get("input_data") if which == "input" else op.get("output_data")
+    shape = op.get("input_shape") if which == "input" else op.get("output_shape")
+    if t is None or shape is None:
+        return jsonify({"error": "Flip tensor data unavailable"}), 200
+    try:
+        if len(shape) <= 1:
+            idx = int(x or 0)
+            return jsonify({"value": t[idx].item()})
+        elif len(shape) == 2:
+            rr = int(y or 0)
+            cc = int(x or 0)
+            return jsonify({"value": t[rr, cc].item()})
+        else:
+            # higher dims not directly supported; fallback to flat index
+            idx = int(x or 0)
+            flat = t.view(-1)
+            return jsonify({"value": flat[idx].item()})
+    except Exception as e:
+        return jsonify({"error": f"Flip get value failed: {e}"}), 200
+
+
 @app.route("/api/getLoadTensor", methods=["POST"])
 def get_load_tensor():
     """Return entire global tensor for a given Load/Store op, with min/max.

From 2b2968d7f71d752313b10f964fc562327bf611c6 Mon Sep 17 00:00:00 2001
From: gujialiang123 <562960306@qq.com>
Date: Fri, 17 Oct 2025 04:03:54 -0400
Subject: [PATCH 45/90] fix(nki): resolve mypy errors (value setter order,
 unary op calls); pass lint

---
 triton_viz/core/nki.py | 24 ++++++++++++------------
 1 file changed, 12 insertions(+), 12 deletions(-)

diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index cb837dad..dc5f1656 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -113,15 +113,15 @@ def dtype(self):
     @property
     def value(self):
         return self._value
-    
-    @property
-    def data(self):
-        return self._value
 
     @value.setter
     def value(self, new_value):
         self._value = new_value
 
+    @property
+    def data(self):
+        return self._value
+
     def data_ptr(self):
         return self._value.ctypes.data
 
@@ -399,16 +399,16 @@ def unary_op(self, x: NDArray, np_func, op_name, **kwargs):
 
     # Elementwise operator implementations
     def exp(self, x: NDArray, **kwargs):
-        return self._unary_op(x, np.exp, "exp", **kwargs)
+        return self.unary_op(x, np.exp, "exp", **kwargs)
 
     def relu(self, x: NDArray, **kwargs):
-        return self._unary_op(x, lambda v: np.maximum(v, 0), "relu", **kwargs)
+        return self.unary_op(x, lambda v: np.maximum(v, 0), "relu", **kwargs)
 
     def sigmoid(self, x: NDArray, **kwargs):
-        return self._unary_op(x, lambda v: 1 / (1 + np.exp(-v)), "sigmoid", **kwargs)
+        return self.unary_op(x, lambda v: 1 / (1 + np.exp(-v)), "sigmoid", **kwargs)
 
     def tanh(self, x: NDArray, **kwargs):
-        return self._unary_op(x, np.tanh, "tanh", **kwargs)
+        return self.unary_op(x, np.tanh, "tanh", **kwargs)
 
     def silu(self, x: NDArray, **kwargs):
         # SiLU(x) = x * sigmoid(x)
@@ -422,13 +422,13 @@ def gelu(self, x: NDArray, **kwargs):
         return NDArray(value=0.5 * x._value * (1 + np.tanh(inner)), name=f"{x.name}_gelu", **kwargs)
 
     def sqrt(self, x: NDArray, **kwargs):
-        return self._unary_op(x, np.sqrt, "sqrt", **kwargs)
+        return self.unary_op(x, np.sqrt, "sqrt", **kwargs)
 
     def abs(self, x: NDArray, **kwargs):
-        return self._unary_op(x, np.abs, "abs", **kwargs)
+        return self.unary_op(x, np.abs, "abs", **kwargs)
 
     def log(self, x: NDArray, **kwargs):
-        return self._unary_op(x, np.log, "log", **kwargs)
+        return self.unary_op(x, np.log, "log", **kwargs)
 
     def pow(self, x: NDArray, exponent, **kwargs):
         if isinstance(exponent, NDArray):
@@ -439,7 +439,7 @@ def pow(self, x: NDArray, exponent, **kwargs):
             raise TypeError(f"Unsupported exponent type: {type(exponent)}")
 
     def reciprocal(self, x: NDArray, **kwargs):
-        return self._unary_op(x, lambda v: 1 / v, "reciprocal", **kwargs)
+        return self.unary_op(x, lambda v: 1 / v, "reciprocal", **kwargs)
 
     def matmul(self, x: NDArray, y: NDArray, transpose_x=False, mask=None, **kwargs):
         x_value = x._value 

From 82f629e665691fbf3c31c69fe9e4a9816cab25fe Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Fri, 24 Oct 2025 19:36:17 -0400
Subject: [PATCH 46/90] use import for nki test

---
 .gitignore        | 2 ++
 tests/test_nki.py | 2 +-
 2 files changed, 3 insertions(+), 1 deletion(-)

diff --git a/.gitignore b/.gitignore
index 060367fa..44a40575 100644
--- a/.gitignore
+++ b/.gitignore
@@ -161,3 +161,5 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+uv.lock
diff --git a/tests/test_nki.py b/tests/test_nki.py
index 3f45c47c..0ec9973d 100644
--- a/tests/test_nki.py
+++ b/tests/test_nki.py
@@ -9,7 +9,7 @@
 import numpy as np
 
 # Import the nki module directly
-exec(open('triton_viz/core/nki.py').read())
+from triton_viz.core.nki import NDArray
 
 def test_ndarray_creation():
     print("Testing NDArray creation...")

From 86345a1f521581fcb9f1e47107fe7c1be30e7b19 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Fri, 24 Oct 2025 22:59:18 -0400
Subject: [PATCH 47/90] masked load fast path; use max dtype value instead of 6
 for out of bounds values

---
 tests/test_masked_load.py          | 140 +++++++++++++++--------------
 triton_viz/core/nki_masked_load.py |  21 ++---
 2 files changed, 80 insertions(+), 81 deletions(-)

diff --git a/tests/test_masked_load.py b/tests/test_masked_load.py
index 40c6323c..c283d1f6 100644
--- a/tests/test_masked_load.py
+++ b/tests/test_masked_load.py
@@ -115,72 +115,73 @@ def test_masked_load():
     mask = np.mgrid[:3, :1, :20, :3, :1, :1]
     mask = (mask[0] > 0) & (mask[2] < 3) & (mask[3] < 2)
     result = masked_load(arr, arr_slice, mask=mask)
+    UD = np.iinfo(arr.dtype).max # UD = undefined (out of bounds)
     expected = np.array([
         [
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
         ],
         [
-            [11, 12,  6],
-            [14, 15,  6],
-            [17, 18,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
+            [11, 12, UD],
+            [14, 15, UD],
+            [17, 18, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
         ],
         [
-            [21, 22,  6],
-            [24, 25,  6],
-            [27, 28,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
+            [21, 22, UD],
+            [24, 25, UD],
+            [27, 28, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
             ],
         ]
     )[:, None, :, :, None, None]
@@ -211,21 +212,22 @@ def test_masked_load():
     mask = np.mgrid[:3, :3, :3]
     mask = (0 < mask[0]) & (mask[0] < 2) & (mask[1] < 2) & (mask[2] != 1)
     result = masked_load(arr, arr_slice, mask=mask)
+    UD = np.iinfo(arr.dtype).max # UD = undefined (out of bounds)
     expected = np.array([
         [
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
         ],
         [
-            [11,  6, 13],
-            [14,  6, 16],
-            [ 6,  6,  6],
+            [11, UD, 13],
+            [14, UD, 16],
+            [UD, UD, UD],
         ],
         [
-            [ 6,  6,  6],
-            [ 6,  6,  6],
-            [ 6,  6,  6],
+            [UD, UD, UD],
+            [UD, UD, UD],
+            [UD, UD, UD],
         ],
     ])
     assert np.allclose(result, expected)
diff --git a/triton_viz/core/nki_masked_load.py b/triton_viz/core/nki_masked_load.py
index 1094756b..a480f006 100644
--- a/triton_viz/core/nki_masked_load.py
+++ b/triton_viz/core/nki_masked_load.py
@@ -73,9 +73,14 @@ def masked_load(ndarray: np.ndarray, keys: tuple, mask: np.ndarray = None) -> np
     Returns:
         Indexed array with masked error handling
     """
-    # Handle mask=None case
-    if mask is None:
-        return ndarray[keys]
+    try: # fast path case - if keys aren't OOB, just go with that
+        out = ndarray[keys]
+        if mask is None:
+            return out
+        out[~mask] = np.iinfo(ndarray.dtype).max
+        return out
+    except:
+        pass
 
     # Convert keys to tuple if it's not already
     if not isinstance(keys, tuple):
@@ -109,7 +114,7 @@ def masked_load(ndarray: np.ndarray, keys: tuple, mask: np.ndarray = None) -> np
         raise IndexError(f"index {oob_idx} is out of bounds for array of size {ndarray.shape}")
 
     valid_mask = mask & in_bounds_mask
-    result[~valid_mask] = 6
+    result[~valid_mask] = np.iinfo(ndarray.dtype).max
 
     return np.expand_dims(result, singleton_dims)
 
@@ -159,12 +164,4 @@ def masked_store(ndarray: np.ndarray, keys: tuple, value: np.ndarray, mask: np.n
         oob_idx = tuple(coord[0] for coord in oob_coords)
         raise IndexError(f"index {oob_idx} is out of bounds for array of size {ndarray.shape}")
 
-    print('!Ndarray')
-    print(ndarray)
-    print('!values')
-    print(value)
-    print('!Offsets')
-    print(offsets)
-    print('!mask')
-    print(mask)
     ndarray.ravel()[offsets[mask]] = value.ravel()[offsets[mask]]

From 85972d8bd19218c86a32f37254f79624d55ec1ab Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Fri, 24 Oct 2025 23:43:13 -0400
Subject: [PATCH 48/90] fix masked store

---
 tests/test_masked_load.py          |  8 ++++----
 triton_viz/core/nki_masked_load.py | 33 ++++++++++++++----------------
 2 files changed, 19 insertions(+), 22 deletions(-)

diff --git a/tests/test_masked_load.py b/tests/test_masked_load.py
index c283d1f6..74596a4e 100644
--- a/tests/test_masked_load.py
+++ b/tests/test_masked_load.py
@@ -291,7 +291,7 @@ def test_masked_load():
         pass
 
 
-    print("Test 10: Make sure error if mask shape wrong")
+    print("Test 11: Make sure error if mask shape wrong")
     arr = np.array([
         [
             [1, 2, 3],
@@ -598,7 +598,7 @@ def test_masked_store():
         print("Success: Correctly raised AssertionError for values/mask shape mismatch\n")
 
     # Test 10: Values shape mismatch
-    print("Test 10: Values shape mismatch should raise AssertionError")
+    print("Test 10: incompatible shape implied by keys and mask")
     arr = np.array([1, 2, 3])
     values = np.array([10, 20, 30, 40])  # Wrong shape
     mask = np.array([True, False, True, False])
@@ -606,8 +606,8 @@ def test_masked_store():
     try:
         masked_store(arr_copy, (slice(0, 3),), values, mask=mask)
         raise RuntimeError("Should've raised an assertion error but did not")
-    except AssertionError:
-        print("Success: Correctly raised AssertionError for values shape mismatch\n")
+    except IndexError:
+        print("Success: Correctly raised IndexError for values shape mismatch\n")
 
     print("All masked_store tests completed!")
 
diff --git a/triton_viz/core/nki_masked_load.py b/triton_viz/core/nki_masked_load.py
index a480f006..089b2b98 100644
--- a/triton_viz/core/nki_masked_load.py
+++ b/triton_viz/core/nki_masked_load.py
@@ -121,6 +121,9 @@ def masked_load(ndarray: np.ndarray, keys: tuple, mask: np.ndarray = None) -> np
 def masked_store(ndarray: np.ndarray, keys: tuple, value: np.ndarray, mask: np.ndarray = None) -> None:
     """
     Store array elements with masking for out-of-bounds errors.
+    General idea of this procedure is to get all the indices that the slice + mask needs for a single advanced indexing assign.
+    - to handle OOB accesses we infer the min possible array size where array[keys] wouldn't have any OOBs
+    - we can get the indices for each dim by selecting from an mgrid
 
     Args:
         ndarray: Input numpy array to store values into
@@ -135,7 +138,6 @@ def masked_store(ndarray: np.ndarray, keys: tuple, value: np.ndarray, mask: np.n
     """
     # Handle mask=None case
     if mask is None:
-        import tp; tp.log(f'{ndarray.shape=}, {keys=}')
         ndarray[keys] = value
         return
 
@@ -145,23 +147,18 @@ def masked_store(ndarray: np.ndarray, keys: tuple, value: np.ndarray, mask: np.n
     if not isinstance(keys, tuple):
         keys = (keys,)
 
+    flat_mask = mask.ravel() # can only index with bool tensors if they're 1d
+
     # normalize (remove Nones in slices)
     keys, singleton_dims = normalize_slice(ndarray, keys)
     target_shape = _calculate_target_shape(keys, ndarray.shape)
-    offset_array = np.arange(np.prod(target_shape)).reshape(target_shape)
-
-    offsets = np.expand_dims(offset_array[keys], singleton_dims)
-    assert offsets.shape == mask.shape
-
-    mask = np.squeeze(mask, singleton_dims)
-
-    # Check if there are any OOB indices where mask=True
-    in_bounds_mask = _get_valid_indices(keys, ndarray.shape, offsets.shape)
-    oob_with_true_mask = (~in_bounds_mask) & mask
-    oob_coords = np.where(oob_with_true_mask)
-    if len(oob_coords[0]) > 0:
-        # Get the first OOB coordinate
-        oob_idx = tuple(coord[0] for coord in oob_coords)
-        raise IndexError(f"index {oob_idx} is out of bounds for array of size {ndarray.shape}")
-
-    ndarray.ravel()[offsets[mask]] = value.ravel()[offsets[mask]]
+    mgrid = np.mgrid[tuple(slice(0, dim, 1) for dim in target_shape)]
+    idxs = []
+    for mgrid_dim in mgrid:
+        # get the section that would be sliced if the array was big enough
+        sliced_section = mgrid_dim[keys]
+
+        # get just the section where mask=True
+        dim_idxs = sliced_section.ravel()[flat_mask]
+        idxs.append(dim_idxs)
+    ndarray[tuple(idxs)] = value.ravel()[flat_mask]

From 30d4e5d156ee3f94c17d5cd02aaeb3d54f965f9c Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Sat, 25 Oct 2025 00:36:22 -0400
Subject: [PATCH 49/90] clone array so load doesn't write on input array

---
 triton_viz/core/nki_masked_load.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/triton_viz/core/nki_masked_load.py b/triton_viz/core/nki_masked_load.py
index 089b2b98..5c5ef184 100644
--- a/triton_viz/core/nki_masked_load.py
+++ b/triton_viz/core/nki_masked_load.py
@@ -74,7 +74,7 @@ def masked_load(ndarray: np.ndarray, keys: tuple, mask: np.ndarray = None) -> np
         Indexed array with masked error handling
     """
     try: # fast path case - if keys aren't OOB, just go with that
-        out = ndarray[keys]
+        out = ndarray[keys].clone()
         if mask is None:
             return out
         out[~mask] = np.iinfo(ndarray.dtype).max

From f6964091fd99375fc9aeab09ca2bea5a167b0e29 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Sat, 25 Oct 2025 00:36:38 -0400
Subject: [PATCH 50/90] make tests prettier

---
 tests/test_masked_load.py | 353 +++++++++++++++++---------------------
 1 file changed, 162 insertions(+), 191 deletions(-)

diff --git a/tests/test_masked_load.py b/tests/test_masked_load.py
index 74596a4e..1fc5ab55 100644
--- a/tests/test_masked_load.py
+++ b/tests/test_masked_load.py
@@ -1,38 +1,49 @@
 import numpy as np
+import pytest
 from triton_viz.core.nki_masked_load import masked_load, masked_store
 
-def test_masked_load():
-    """Test the masked_load function with various scenarios."""
-
-    print("=== Testing masked_load function ===\n")
-
-    # Test 1: mask=None case (direct indexing)
-    print("Test 1: mask=None case")
+def print_op_details(test_name, operation_type, input_data, keys, values=None, mask=None, output=None, error_expected=None):
+    """Print detailed information about a masked load/store operation."""
+    print(f"{test_name}:")
+    print(f"Operation type: {operation_type}")
+    print("---------------------------------")
+    print(f"Input:\n{input_data}")
+    print("---------------------------------")
+    print(f"Values:\n{values}")
+    print("---------------------------------")
+    print(f"Keys:\n{keys}")
+    print("---------------------------------")
+    print(f"Mask:\n{mask}")
+    print("---------------------------------")
+    print(f"Result {operation_type}:\n{output}")
+    if error_expected:
+        print(f"Expecting {error_expected}...")
+    print()
+
+def test_masked_load_none_case():
+    """Test masked_load with mask=None (direct indexing)."""
     arr = np.array([1, 2, 3, 4, 5])
     result = masked_load(arr, (slice(1, 4),), mask=None)
     expected = arr[1:4]
-    print(f"Input: {arr}")
-    print(f"Keys: slice(1, 4)")
-    print(f"Result: {result}")
-    print(f"Expected: {expected}")
-    print(f"Match: {np.array_equal(result, expected)}\n")
-
-    # Test 2: In-bounds indexing with mask
-    print("Test 2: In-bounds indexing with mask")
+
+    print_op_details("Test 1: mask=None case", "load", arr, (slice(1, 4),), mask=None, output=result)
+
+    assert np.array_equal(result, expected)
+
+def test_masked_load_in_bounds():
+    """Test masked_load with in-bounds indexing and mask."""
     arr = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
     mask = np.array([[True, True, False], [True, False, True], [False, True, True]])
     result = masked_load(arr, (slice(0, 3), slice(0, 3)), mask=mask)
-    result[[0,1,2],[2,1,0]] = 0
+    result[[0,1,2],[2,1,0]] = 0  # manually set masked out positions to 0 for testing
     expected = np.array([[1, 2, 0], [4, 0, 6], [0, 8, 9]])
-    print(f"Input array:\n{arr}")
-    print(f"Mask:\n{mask}")
-    print(f"Keys: (slice(0, 2), slice(0, 2))")
-    print(f"Result:\n{result}")
-    print(f"Expected:\n{expected}")
-    print(f"Match: {np.array_equal(result, expected)}\n")
 
-    # Test 3: Out-of-bounds indexing with mask
-    print("Test 3: Out-of-bounds indexing with mask")
+    print_op_details("Test 2: In-bounds indexing with mask", "load", arr, (slice(0, 3), slice(0, 3)), mask=mask, output=result)
+
+    assert np.array_equal(result, expected)
+
+def test_masked_load_out_of_bounds():
+    """Test masked_load with out-of-bounds indexing and mask."""
     arr = np.array([[1, 2], [3, 4]])
     mask = np.array([
         [True , False, False],
@@ -40,53 +51,49 @@ def test_masked_load():
         [False, False, False],
         [False, False, False],
         [False, False, False],
-    ])  # Shape doesn't match exactly, but demonstrates concept
-    # This would normally cause IndexError for arr[5], but with mask should work
+    ])  # note that arr.shape != mask.shape, this is intended
+
+    # normally this arr[slice] would OOB but the mask at the OOB idxs are false so they're not loaded
     result = masked_load(arr, (slice(0, 5), slice(0, 3)), mask=mask)
-    print(f"Input array:\n{arr}")
-    print(f"Keys: (slice(0, 5), slice(0, 3)) - exceeds array bounds")
-    print(f"Result shape: {result.shape}")
-    print(f"Result:\n{result}")
-    print("Success: Handled OOB indexing\n")
-
-    # Test 4: Integer indexing OOB
-    print("Test 4: Integer indexing out of bounds")
+
+    print_op_details("Test 3: Out-of-bounds indexing with mask", "load", arr, (slice(0, 5), slice(0, 3)), mask=mask, output=result)
+
+    assert result.shape == (5, 3)
+
+def test_masked_load_integer_oob():
+    """Test masked_load with integer out-of-bounds indexing."""
     arr = np.array([10, 20, 30])
-    mask = np.array([False])  # Extra mask element
-    result = masked_load(arr, ([5],), mask=mask)  # Index 5 is OOB
-    print(f"Input: {arr}")
-    print(f"Keys: (5,) - OOB index")
-    print(f"Result: {result}")
-    print(f"Result type: {type(result)}")
-    print("Success: Handled OOB integer indexing\n")
-
-    # Test 5: Array indexing
-    print("Test 5: Array indexing")
+    mask = np.array([False])
+    result = masked_load(arr, ([5],), mask=mask)  # idx 5 is OOB
+
+    print_op_details("Test 4: Integer indexing out of bounds", "load", arr, ([5],), mask=mask, output=result)
+
+    assert isinstance(result, np.ndarray)
+
+def test_masked_load_array_indexing():
+    """Test masked_load with array indexing."""
     arr = np.array([100, 200, 300, 400, 500])
-    indices = np.array([0, 2, 4])  # Valid indices
+    indices = np.array([0, 2, 4])
     mask = np.array([True, True, True])
     result = masked_load(arr, (indices,), mask=mask)
     expected = arr[indices]
-    print(f"Input: {arr}")
-    print(f"Indices: {indices}")
-    print(f"Result: {result}")
-    print(f"Expected: {expected}")
-    print(f"Match: {np.array_equal(result, expected)}\n")
-
-    # Test 6: Mixed indexing types
-    print("Test 6: Mixed indexing types")
+
+    print_op_details("Test 5: Array indexing", "load", arr, (indices,), mask=mask, output=result)
+
+    assert np.array_equal(result, expected)
+
+def test_masked_load_mixed_indexing():
+    """Test masked_load with mixed indexing types."""
     arr = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])
     mask = np.array([True, False, True, False])
-    # Mix of slice and integer
-    result = masked_load(arr, (slice(0, 4), 2), mask=mask)  # Get column 2, allow extending rows
-    print(f"Input array:\n{arr}")
-    print(f"Keys: (slice(0, 4), 2) - slice extends beyond array rows")
-    print(f"Result: {result}")
-    print(f"Result shape: {result.shape}")
-    print("Success: Handled mixed indexing\n")
-
-    # Test 7: Mixed indexing types
-    print("Test 7: Mixed indexing types")
+    result = masked_load(arr, (slice(0, 4), 2), mask=mask)  # set col 2
+
+    print_op_details("Test 6: Mixed indexing types", "load", arr, (slice(0, 4), 2), mask=mask, output=result)
+
+    assert result.shape == (4,)
+
+def test_masked_load_complex_indexing():
+    """Test masked_load with complex multi-dimensional indexing."""
     arr = np.array([
         [
             [1, 2, 3],
@@ -115,7 +122,7 @@ def test_masked_load():
     mask = np.mgrid[:3, :1, :20, :3, :1, :1]
     mask = (mask[0] > 0) & (mask[2] < 3) & (mask[3] < 2)
     result = masked_load(arr, arr_slice, mask=mask)
-    UD = np.iinfo(arr.dtype).max # UD = undefined (out of bounds)
+    UD = np.iinfo(arr.dtype).max  # UD = undefined (out of bounds)
     expected = np.array([
         [
             [UD, UD, UD],
@@ -185,10 +192,13 @@ def test_masked_load():
             ],
         ]
     )[:, None, :, :, None, None]
+
+    print("Load Test 7: Complex multi-dimensional indexing (omitted details, check code to see inputs/outputs)")
+
     assert np.allclose(result, expected)
 
-    # Test 8: incomplete slices
-    print("Test 8: incomplete slices")
+def test_masked_load_incomplete_slices():
+    """Test masked_load with incomplete slices."""
     arr = np.array([
         [
             [1, 2, 3],
@@ -206,13 +216,11 @@ def test_masked_load():
             [27, 28, 29],
         ],
     ])
-    arr_slice = (
-        slice(None, None, None),
-    )
+    arr_slice = (slice(None, None, None),)
     mask = np.mgrid[:3, :3, :3]
     mask = (0 < mask[0]) & (mask[0] < 2) & (mask[1] < 2) & (mask[2] != 1)
     result = masked_load(arr, arr_slice, mask=mask)
-    UD = np.iinfo(arr.dtype).max # UD = undefined (out of bounds)
+    UD = np.iinfo(arr.dtype).max  # UD = undefined (out of bounds)
     expected = np.array([
         [
             [UD, UD, UD],
@@ -230,9 +238,13 @@ def test_masked_load():
             [UD, UD, UD],
         ],
     ])
+
+    print_op_details("Test 8: Incomplete slices", "load", arr, arr_slice, mask=mask, output=result)
+
     assert np.allclose(result, expected)
 
-    print("Test 9: Make sure index error still happen when mask=True")
+def test_masked_load_index_error_with_true_mask():
+    """Test that IndexError is raised when mask=True at out-of-bounds location."""
     arr = np.array([
         [
             [1, 2, 3],
@@ -250,20 +262,19 @@ def test_masked_load():
             [27, 28, 29],
         ],
     ])
-    arr_slice = (
-        slice(None, None, None),
-        slice(0, 4)
-    )
+    arr_slice = (slice(None, None, None), slice(0, 4))
     mask = np.mgrid[:3, :4, :3]
     mask = (mask[0] < 3) & (mask[1] < 3) & (mask[2] < 3)
-    mask[-1, -1, -1] = True # index error
-    try:
-        result = masked_load(arr, arr_slice, mask=mask)
-        raise RuntimeError("Should've raised an index error but did not")
-    except IndexError:
-        pass
-
-    print("Test 10: Make sure error if mask shape wrong")
+    mask[-1, -1, -1] = True  # index error
+
+    print_op_details("Test 9: Index error should still happen when mask=True at OOB location", "load", arr, arr_slice, mask=mask, error_expected="IndexError")
+
+    with pytest.raises(IndexError):
+        masked_load(arr, arr_slice, mask=mask)
+    print("Success: Correctly raised IndexError\n")
+
+def test_masked_load_mask_shape_mismatch():
+    """Test that AssertionError is raised when mask shape doesn't match array slice."""
     arr = np.array([
         [
             [1, 2, 3],
@@ -284,14 +295,15 @@ def test_masked_load():
     arr_slice = (slice(None, None, None),)
     mask = np.mgrid[:3, :4, :3]
     mask = (mask[0] < 3) & (mask[1] < 3) & (mask[2] < 3)
-    try:
-        result = masked_load(arr, arr_slice, mask=mask)
-        raise RuntimeError("Should've raised an assertion error (arr[arr_slice].shape != mask.shape) but did not")
-    except AssertionError:
-        pass
 
+    print_op_details("Test 10: Error if mask shape wrong", "load", arr, arr_slice, mask=mask, error_expected="AssertionError")
 
-    print("Test 11: Make sure error if mask shape wrong")
+    with pytest.raises(AssertionError):
+        masked_load(arr, arr_slice, mask=mask)
+    print("Success: Correctly raised AssertionError\n")
+
+def test_masked_load_correct_mask_shape():
+    """Test masked_load with correctly shaped mask."""
     arr = np.array([
         [
             [1, 2, 3],
@@ -309,47 +321,45 @@ def test_masked_load():
             [27, 28, 29],
         ],
     ])
+    print(arr)
     arr_slice = (slice(None, None, None),)
     mask = np.arange(27).reshape(3, 3, 3) % 2 == 0
     result = masked_load(arr, arr_slice, mask=mask)
-    print(result)
 
-def test_masked_store():
-    """Test the masked_store function with various scenarios."""
+    print(arr)
+    raise
+    print_op_details("Test 11: Correct mask shape", "load", arr, arr_slice, mask=mask, output=result)
 
-    print("=== Testing masked_store function ===\n")
+    assert result.shape == mask.shape
 
-    # Test 1: mask=None case (direct indexing)
-    print("Test 1: mask=None case")
+def test_masked_store_none_case():
+    """Test masked_store with mask=None (direct indexing)."""
     arr = np.array([1, 2, 3, 4, 5])
     values = np.array([10, 20, 30])
     arr_copy = arr.copy()
     masked_store(arr_copy, (slice(1, 4),), values, mask=None)
     expected = arr.copy()
     expected[1:4] = values
-    print(f"Original: {arr}")
-    print(f"Values: {values}")
-    print(f"Result: {arr_copy}")
-    print(f"Expected: {expected}")
+
+    print_op_details("Store Test 1: mask=None case", "store", arr, (slice(1, 4),), values=values, mask=None, output=arr_copy)
+
     assert np.array_equal(arr_copy, expected)
 
-    # Test 2: In-bounds indexing with mask
-    print("Test 2: In-bounds indexing with mask")
+def test_masked_store_in_bounds():
+    """Test masked_store with in-bounds indexing and mask."""
     arr = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
     values = np.array([[100, 200, 0], [400, 0, 600], [0, 800, 900]])
     mask = np.array([[True, True, False], [True, False, True], [False, True, True]])
     arr_copy = arr.copy()
     masked_store(arr_copy, (slice(0, 3), slice(0, 3)), values, mask=mask)
     expected = np.array([[100, 200, 3], [400, 5, 600], [7, 800, 900]])
-    print(f"Original array:\n{arr}")
-    print(f"Values:\n{values}")
-    print(f"Mask:\n{mask}")
-    print(f"Result:\n{arr_copy}")
-    print(f"Expected:\n{expected}")
+
+    print_op_details("Store Test 2: In-bounds indexing with mask", "store", arr, (slice(0, 3), slice(0, 3)), values=values, mask=mask, output=arr_copy)
+
     assert np.array_equal(arr_copy, expected)
 
-    # Test 3: Out-of-bounds indexing with mask (should succeed where mask=False)
-    print("Test 3: Out-of-bounds indexing with mask")
+def test_masked_store_out_of_bounds():
+    """Test masked_store with out-of-bounds indexing and mask."""
     arr = np.array([[1, 2], [3, 4]])
     values = np.array([
         [100, 0, 0],
@@ -368,27 +378,25 @@ def test_masked_store():
     arr_copy = arr.copy()
     masked_store(arr_copy, (slice(0, 5), slice(0, 3)), values, mask=mask)
     expected = np.array([[100, 2], [3, 200]])
-    print(f"Original array:\n{arr}")
-    print(f"Keys: (slice(0, 5), slice(0, 3)) - exceeds array bounds")
-    print(f"Result:\n{arr_copy}")
-    print(f"Expected:\n{expected}")
+
+    print_op_details("Store Test 3: Out-of-bounds indexing with mask", "store", arr, (slice(0, 5), slice(0, 3)), values=values, mask=mask, output=arr_copy)
+
     assert np.array_equal(arr_copy, expected)
 
-    # Test 4: Integer indexing OOB
-    print("Test 4: Integer indexing out of bounds")
+def test_masked_store_integer_oob():
+    """Test masked_store with integer out-of-bounds indexing."""
     arr = np.array([10, 20, 30])
     values = np.array([999])
     mask = np.array([False])  # mask=False at OOB index
     arr_copy = arr.copy()
     masked_store(arr_copy, ([5],), values, mask=mask)  # Index 5 is OOB
-    print(f"Original: {arr}")
-    print(f"Keys: (5,) - OOB index")
-    print(f"Result: {arr_copy}")
-    print(f"Expected: {arr} (unchanged)")
+
+    print_op_details("Store Test 4: Integer indexing out of bounds", "store", arr, ([5],), values=values, mask=mask, output=arr_copy)
+
     assert np.array_equal(arr_copy, arr)
 
-    # Test 5: Array indexing
-    print("Test 5: Array indexing")
+def test_masked_store_array_indexing():
+    """Test masked_store with array indexing."""
     arr = np.array([100, 200, 300, 400, 500])
     indices = np.array([0, 2, 4])
     values = np.array([111, 333, 555])
@@ -396,30 +404,26 @@ def test_masked_store():
     arr_copy = arr.copy()
     masked_store(arr_copy, (indices,), values, mask=mask)
     expected = np.array([111, 200, 333, 400, 555])
-    print(f"Original: {arr}")
-    print(f"Indices: {indices}")
-    print(f"Values: {values}")
-    print(f"Result: {arr_copy}")
-    print(f"Expected: {expected}")
+
+    print_op_details("Store Test 5: Array indexing", "store", arr, (indices,), values=values, mask=mask, output=arr_copy)
+
     assert np.array_equal(arr_copy, expected)
 
-    # Test 6: Mixed indexing types
-    print("Test 6: Mixed indexing types")
+def test_masked_store_mixed_indexing():
+    """Test masked_store with mixed indexing types."""
     arr = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])
     values = np.array([100, 200, 300, 400])
     mask = np.array([True, False, True, False])
     arr_copy = arr.copy()
-    masked_store(arr_copy, (slice(0, 4), 2), values, mask=mask)  # Set column 2
+    masked_store(arr_copy, (slice(0, 4), 2), values, mask=mask)  # set col 2
     expected = np.array([[1, 2, 100, 4], [5, 6, 7, 8], [9, 10, 300, 12]])
-    print(f"Original array:\n{arr}")
-    print(f"Values: {values}")
-    print(f"Keys: (slice(0, 4), 2) - slice extends beyond array rows")
-    print(f"Result:\n{arr_copy}")
-    print(f"Expected:\n{expected}")
+
+    print_op_details("Store Test 6: Mixed indexing types", "store", arr, (slice(0, 4), 2), values=values, mask=mask, output=arr_copy)
+
     assert np.array_equal(arr_copy, expected)
 
-    # Test 7: Complex 3D indexing (simplified)
-    print("Test 7: Complex 3D indexing")
+def test_masked_store_complex_indexing():
+    """Test masked_store with complex multi-dimensional indexing."""
     arr = np.array([
         [
             [1, 2, 3],
@@ -437,28 +441,6 @@ def test_masked_store():
             [27, 28, 29],
         ],
     ])
-    ## Simpler test case that uses basic slicing with mask
-    #arr_slice = (slice(0, 3), slice(0, 3))  # Just use 2D slicing for simplicity
-    #mask = np.array([
-    #    [True, True, False],
-    #    [False, True, True],
-    #    [True, False, False]
-    #])
-    #values = np.array([
-    #    [100, 200, 0],
-    #    [0, 300, 400],
-    #    [500, 0, 0]
-    #])
-    #arr_copy = arr[:, :, 0]  # Take first slice along last dimension
-    #arr_copy = arr_copy.copy()
-    #masked_store(arr_copy, arr_slice, values, mask=mask)
-
-    #expected_modified = np.array([
-    #    [100, 200, 7],   # Original row 0: [1, 4, 7], mask: [True, True, False] -> [100, 200, 7]
-    #    [11, 300, 400],  # Original row 1: [11, 14, 17], mask: [False, True, True] -> [11, 300, 400]
-    #    [500, 24, 27]    # Original row 2: [21, 24, 27], mask: [True, False, False] -> [500, 24, 27]
-    #])
-
     arr_slice = (
         slice(None, None, None),
         None,
@@ -469,7 +451,6 @@ def test_masked_store():
     )
     mask = np.mgrid[:3, :1, :20, :3, :1, :1]
     mask = (mask[0] > 0) & (mask[2] < 3) & (mask[3] < 2)
-    result = masked_load(arr, arr_slice, mask=mask)
     values = np.array([
         [
             [99, 99, 99],
@@ -541,7 +522,6 @@ def test_masked_store():
     )[:, None, :, :, None, None]
     arr_copy = arr.copy()
     masked_store(arr_copy, arr_slice, values, mask=mask)
-
     expected_modified = np.array([
         [
             [1, 2, 3],
@@ -560,14 +540,12 @@ def test_masked_store():
         ],
     ])
 
-    print(f"Input array shape: {arr_copy.shape}")
-    print(f"Keys: 2D slicing")
-    print(f"Values shape: {values.shape}")
-    print(f"Mask shape: {mask.shape}")
+    print("Store Test 7: Complex multi-dimensional indexing (omitted details, check code to see inputs/outputs)")
+
     assert np.array_equal(arr_copy, expected_modified)
 
-    # Test 8: Error case - OOB with mask=True should raise IndexError
-    print("Test 8: OOB with mask=True should raise IndexError")
+def test_masked_store_oob_with_true_mask():
+    """Test that IndexError is raised when mask=True at out-of-bounds location."""
     arr = np.array([[1, 2], [3, 4]])
     values = np.array([[100, 200], [300, 400], [500, 600]])
     mask = np.array([
@@ -576,42 +554,35 @@ def test_masked_store():
         [True, False],  # This position is OOB and mask=True
     ])
     arr_copy = arr.copy()
-    try:
+
+    print_op_details("Store Test 8: OOB with mask=True should raise IndexError", "store", arr, (slice(0, 3), slice(0, 2)), values=values, mask=mask, error_expected="IndexError")
+
+    with pytest.raises(IndexError):
         masked_store(arr_copy, (slice(0, 3), slice(0, 2)), values, mask=mask)
-        raise RuntimeError("Should've raised an index error but did not")
-    except IndexError:
-        print("Success: Correctly raised IndexError for OOB with mask=True\n")
+    print("Success: Correctly raised IndexError\n")
 
-    # Test 9: Error case - values and mask shape mismatch
-    print("Test 9: Values and mask shape mismatch should raise AssertionError")
+def test_masked_store_values_mask_shape_mismatch():
+    """Test that AssertionError is raised when values and mask shapes don't match."""
     arr = np.array([1, 2, 3])
     values = np.array([10, 20, 30])  # Shape (3,)
     mask = np.array([True, False])   # Shape (2,) - different from values
     arr_copy = arr.copy()
-    try:
+
+    print_op_details("Store Test 9: Values and mask shape mismatch should raise AssertionError", "store", arr, (slice(0, 3),), values=values, mask=mask, error_expected="AssertionError")
+
+    with pytest.raises(AssertionError):
         masked_store(arr_copy, (slice(0, 3),), values, mask=mask)
-        print(f"ERROR: Should have raised AssertionError but didn't")
-        print(f"values.shape: {values.shape}")
-        print(f"mask.shape: {mask.shape}")
-        raise RuntimeError("Should've raised an assertion error but did not")
-    except AssertionError:
-        print("Success: Correctly raised AssertionError for values/mask shape mismatch\n")
-
-    # Test 10: Values shape mismatch
-    print("Test 10: incompatible shape implied by keys and mask")
+    print("Success: Correctly raised AssertionError\n")
+
+def test_masked_store_values_shape_mismatch():
+    """Test that IndexError is raised when values shape is incompatible."""
     arr = np.array([1, 2, 3])
     values = np.array([10, 20, 30, 40])  # Wrong shape
     mask = np.array([True, False, True, False])
     arr_copy = arr.copy()
-    try:
-        masked_store(arr_copy, (slice(0, 3),), values, mask=mask)
-        raise RuntimeError("Should've raised an assertion error but did not")
-    except IndexError:
-        print("Success: Correctly raised IndexError for values shape mismatch\n")
 
-    print("All masked_store tests completed!")
+    print_op_details("Store Test 10: Values shape mismatch", "store", arr, (slice(0, 3),), values=values, mask=mask, error_expected="IndexError")
 
-if __name__ == "__main__":
-    test_masked_load()
-    print("\n" + "="*50 + "\n")
-    test_masked_store()
+    with pytest.raises(IndexError):
+        masked_store(arr_copy, (slice(0, 3),), values, mask=mask)
+    print("Success: Correctly raised IndexError\n")

From 9d9649960da0c57ae64ff056d2c3206870081f8b Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Sat, 25 Oct 2025 00:57:16 -0400
Subject: [PATCH 51/90] make all tests pass

---
 tests/test_masked_load.py          | 3 ---
 triton_viz/core/nki_masked_load.py | 2 +-
 2 files changed, 1 insertion(+), 4 deletions(-)

diff --git a/tests/test_masked_load.py b/tests/test_masked_load.py
index 1fc5ab55..5eedc3cf 100644
--- a/tests/test_masked_load.py
+++ b/tests/test_masked_load.py
@@ -321,13 +321,10 @@ def test_masked_load_correct_mask_shape():
             [27, 28, 29],
         ],
     ])
-    print(arr)
     arr_slice = (slice(None, None, None),)
     mask = np.arange(27).reshape(3, 3, 3) % 2 == 0
     result = masked_load(arr, arr_slice, mask=mask)
 
-    print(arr)
-    raise
     print_op_details("Test 11: Correct mask shape", "load", arr, arr_slice, mask=mask, output=result)
 
     assert result.shape == mask.shape
diff --git a/triton_viz/core/nki_masked_load.py b/triton_viz/core/nki_masked_load.py
index 5c5ef184..e253c791 100644
--- a/triton_viz/core/nki_masked_load.py
+++ b/triton_viz/core/nki_masked_load.py
@@ -74,7 +74,7 @@ def masked_load(ndarray: np.ndarray, keys: tuple, mask: np.ndarray = None) -> np
         Indexed array with masked error handling
     """
     try: # fast path case - if keys aren't OOB, just go with that
-        out = ndarray[keys].clone()
+        out = ndarray[keys].copy()
         if mask is None:
             return out
         out[~mask] = np.iinfo(ndarray.dtype).max

From 9066d3be9ed9debf9c27ffb12edf867539866a1d Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Sat, 25 Oct 2025 08:46:16 -0400
Subject: [PATCH 52/90] modify undefined values for other dtypes in masked
 load; lint

---
 triton_viz/core/nki_masked_load.py | 43 ++++++++++++++++++++++--------
 1 file changed, 32 insertions(+), 11 deletions(-)

diff --git a/triton_viz/core/nki_masked_load.py b/triton_viz/core/nki_masked_load.py
index e253c791..f13ff19a 100644
--- a/triton_viz/core/nki_masked_load.py
+++ b/triton_viz/core/nki_masked_load.py
@@ -2,8 +2,9 @@
 
 get = lambda x, default: default if x is None else x
 
+
 def normalize_slice(ndarray: np.ndarray, keys: tuple) -> tuple[tuple, tuple]:
-    """ Separate singleton dims (None) and add bounds to slices like :2, 2:, or :"""
+    """Separate singleton dims (None) and add bounds to slices like :2, 2:, or :"""
     singleton_dims = []
     dim_idx = 0
     new_keys = []
@@ -28,7 +29,10 @@ def normalize_slice(ndarray: np.ndarray, keys: tuple) -> tuple[tuple, tuple]:
         dim_idx += 1
     return tuple(new_keys), tuple(singleton_dims)
 
-def _calculate_target_shape(keys: tuple, original_shape: tuple[int, ...]) -> tuple[int, ...]:
+
+def _calculate_target_shape(
+    keys: tuple, original_shape: tuple[int, ...]
+) -> tuple[int, ...]:
     """Calculate the shape needed for the padded array to handle indexing."""
     target_shape = list(original_shape)
     for i, (arr_dim, key) in enumerate(zip(original_shape, keys)):
@@ -40,7 +44,10 @@ def _calculate_target_shape(keys: tuple, original_shape: tuple[int, ...]) -> tup
             target_shape[i] = max(arr_dim, np.max(key) + 1)
     return tuple(target_shape)
 
-def _get_valid_indices(keys: tuple, original_shape: tuple[int, ...], result_shape: tuple[int, ...]) -> np.ndarray:
+
+def _get_valid_indices(
+    keys: tuple, original_shape: tuple[int, ...], result_shape: tuple[int, ...]
+) -> np.ndarray:
     """Get a boolean mask indicating which result indices are within original array bounds."""
     coords = np.mgrid[[slice(0, s) for s in result_shape]]
     valid_mask = np.ones(result_shape, dtype=bool)
@@ -55,11 +62,14 @@ def _get_valid_indices(keys: tuple, original_shape: tuple[int, ...], result_shap
             valid_indices = np.array(key)
             valid_mask &= (0 <= valid_indices) & (valid_indices < arr_dim)
         elif isinstance(key, (int, np.integer)):
-            valid_mask &= (0 <= int(key) < arr_dim)
+            valid_mask &= 0 <= int(key) < arr_dim
 
     return valid_mask
 
-def masked_load(ndarray: np.ndarray, keys: tuple, mask: np.ndarray = None) -> np.ndarray:
+
+def masked_load(
+    ndarray: np.ndarray, keys: tuple, mask: np.ndarray = None
+) -> np.ndarray:
     """
     Load array elements with masking for out-of-bounds errors.
 
@@ -73,13 +83,13 @@ def masked_load(ndarray: np.ndarray, keys: tuple, mask: np.ndarray = None) -> np
     Returns:
         Indexed array with masked error handling
     """
-    try: # fast path case - if keys aren't OOB, just go with that
+    try:  # fast path case - if keys aren't OOB, just go with that
         out = ndarray[keys].copy()
         if mask is None:
             return out
         out[~mask] = np.iinfo(ndarray.dtype).max
         return out
-    except:
+    except Exception:
         pass
 
     # Convert keys to tuple if it's not already
@@ -111,14 +121,25 @@ def masked_load(ndarray: np.ndarray, keys: tuple, mask: np.ndarray = None) -> np
     if len(oob_coords[0]) > 0:
         # Get the first OOB coordinate
         oob_idx = tuple(coord[0] for coord in oob_coords)
-        raise IndexError(f"index {oob_idx} is out of bounds for array of size {ndarray.shape}")
+        raise IndexError(
+            f"index {oob_idx} is out of bounds for array of size {ndarray.shape}"
+        )
 
     valid_mask = mask & in_bounds_mask
-    result[~valid_mask] = np.iinfo(ndarray.dtype).max
+    # Use appropriate max value based on dtype
+    if np.issubdtype(ndarray.dtype, np.integer):
+        result[~valid_mask] = np.iinfo(ndarray.dtype).max
+    elif np.issubdtype(ndarray.dtype, np.floating):
+        result[~valid_mask] = np.finfo(ndarray.dtype).max
+    else:
+        result[~valid_mask] = 0  # false
 
     return np.expand_dims(result, singleton_dims)
 
-def masked_store(ndarray: np.ndarray, keys: tuple, value: np.ndarray, mask: np.ndarray = None) -> None:
+
+def masked_store(
+    ndarray: np.ndarray, keys: tuple, value: np.ndarray, mask: np.ndarray = None
+) -> None:
     """
     Store array elements with masking for out-of-bounds errors.
     General idea of this procedure is to get all the indices that the slice + mask needs for a single advanced indexing assign.
@@ -147,7 +168,7 @@ def masked_store(ndarray: np.ndarray, keys: tuple, value: np.ndarray, mask: np.n
     if not isinstance(keys, tuple):
         keys = (keys,)
 
-    flat_mask = mask.ravel() # can only index with bool tensors if they're 1d
+    flat_mask = mask.ravel()  # can only index with bool tensors if they're 1d
 
     # normalize (remove Nones in slices)
     keys, singleton_dims = normalize_slice(ndarray, keys)

From 696587d1c3b97d7895ed33a76df0b24899b6de02 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Sat, 25 Oct 2025 09:13:09 -0400
Subject: [PATCH 53/90] remove ._value, just use .data; lint

---
 nki-examples/matmul.py |  94 ++++++++--------
 tests/test_nki.py      |  81 +++++++-------
 triton_viz/core/nki.py | 241 ++++++++++++++++++++++-------------------
 3 files changed, 217 insertions(+), 199 deletions(-)

diff --git a/nki-examples/matmul.py b/nki-examples/matmul.py
index 4e3c458b..2a8b8edc 100644
--- a/nki-examples/matmul.py
+++ b/nki-examples/matmul.py
@@ -1,18 +1,15 @@
 from neuronxcc import nki
 import neuronxcc.nki.language as nl
 
-import tp
-import torch
 import triton_viz
-import torch_xla.core.xla_model as xm
 from triton_viz.clients import Tracer
-from triton_viz.core import config as cfg
 from triton_viz.core.trace import launches
 import numpy as np
-import math
 
-tensors = [] # THTODO: rm
-records = [triton_viz.core.data.Grid((0,0,0))] # THTODO: rm
+tensors = set()  # THTODO: rm
+records = [triton_viz.core.data.Grid((0, 0, 0))]  # THTODO: rm
+
+
 def matmul_kernel(lhs, rhs):
     """NKI matmul_kernel to compute a matrix multiplication operation in a tiled manner
 
@@ -36,12 +33,10 @@ def matmul_kernel(lhs, rhs):
     TILE_K = 2
     TILE_N = 4
 
-    tensors.append(lhs)
-    tensors.append(rhs)
-    tensors.append(result)
+    tensors.add(lhs)
+    tensors.add(rhs)
+    tensors.add(result)
     # Use affine_range to loop over tiles
-    m_range = nl.arange(M)[:, None]
-    n_range = nl.arange(N)[None, :]
     for m in nl.affine_range(M // TILE_M):
         for n in nl.affine_range(N // TILE_N):
             # Allocate a tensor in PSUM
@@ -57,33 +52,35 @@ def matmul_kernel(lhs, rhs):
                 lhs_f = nl.arange(TILE_K)[None, :] + k * TILE_K
                 lhs_mask = (lhs_p < M) & (lhs_f < K)
                 lhs_tile = nl.load(lhs[lhs_p, lhs_f], mask=lhs_mask)
-                records.append(triton_viz.core.data.Load(
-                    lhs.data_ptr(),
-                    (4 * (lhs_p * K + lhs_f))._value,
-                    lhs_mask._value
-                )) # THTODO: rm
+                records.append(
+                    triton_viz.core.data.Load(
+                        lhs.data_ptr(), (4 * (lhs_p * K + lhs_f)).data, lhs_mask.data
+                    )
+                )  # THTODO: rm
 
                 rhs_p = nl.arange(TILE_K)[:, None] + k * TILE_K
                 rhs_f = nl.arange(TILE_N)[None, :] + n * TILE_N
-                rhs_mask = (rhs_p < K) & (rhs_f < N) 
+                rhs_mask = (rhs_p < K) & (rhs_f < N)
                 rhs_tile = nl.load(rhs[rhs_p, rhs_f], mask=rhs_mask)
-                records.append(triton_viz.core.data.Load(
-                    rhs.data_ptr(),
-                    (4 * (rhs_p * N + rhs_f))._value,
-                    rhs_mask._value
-                )) # THTODO: rm
+                records.append(
+                    triton_viz.core.data.Load(
+                        rhs.data_ptr(), (4 * (rhs_p * N + rhs_f)).data, rhs_mask.data
+                    )
+                )  # THTODO: rm
 
                 # Accumulate partial-sums into PSUM
                 x = nl.matmul(lhs_tile[...], rhs_tile[...], transpose_x=False)
                 res_psum += x
-                records.append(triton_viz.core.data.Dot(
-                    lhs_tile.shape,
-                    rhs_tile.shape,
-                    res_psum.shape,
-                    lhs_tile._value,
-                    rhs_tile._value,
-                    x._value
-                )) # THTODO: rm
+                records.append(
+                    triton_viz.core.data.Dot(
+                        lhs_tile.shape,
+                        rhs_tile.shape,
+                        res_psum.shape,
+                        lhs_tile.data,
+                        rhs_tile.data,
+                        x.data,
+                    )
+                )  # THTODO: rm
 
             # Copy the result from PSUM back to SBUF, and cast to expected output data-type
             res_sb = nl.copy(res_psum, dtype=result.dtype)
@@ -91,16 +88,19 @@ def matmul_kernel(lhs, rhs):
             out_p = nl.arange(TILE_M)[:, None] + m * TILE_M
             out_f = nl.arange(TILE_N)[None, :] + n * TILE_N
             out_mask = (out_p < M) & (out_f < N)
-            nl.store(result[m * TILE_M:(m + 1) * TILE_M, n * TILE_N:(n + 1) * TILE_N],
-                     value=res_sb)
-            records.append(triton_viz.core.data.Store(
-                result.data_ptr(),
-                (4 * (out_p * N + out_f))._value,
-                out_mask._value
-            )) # THTODO: rm
+            nl.store(
+                result[m * TILE_M : (m + 1) * TILE_M, n * TILE_N : (n + 1) * TILE_N],
+                value=res_sb,
+            )
+            records.append(
+                triton_viz.core.data.Store(
+                    result.data_ptr(), (4 * (out_p * N + out_f)).data, out_mask.data
+                )
+            )  # THTODO: rm
 
     return result
 
+
 TRITON_VIZ = True
 kernel_grid = (1, 1, 1)
 lhs_small = np.arange(16).astype(np.float32).reshape(4, 4)
@@ -108,16 +108,16 @@ def matmul_kernel(lhs, rhs):
 kernel_args = (lhs_small, rhs_small)
 
 if TRITON_VIZ:
-    print('Executing matmul_kernel with NKI interpreter...')
-    matmul_kernel = triton_viz.trace(clients=Tracer(), backend="nki")(matmul_kernel)
-    kk = matmul_kernel[kernel_grid]
+    print("Executing matmul_kernel with NKI interpreter...")
+    traced_kernel = triton_viz.trace(clients=Tracer(), backend="nki")(matmul_kernel)
+    kk = traced_kernel[kernel_grid]
     z2 = kk(*kernel_args)
 
     print(f"Number of launches: {len(launches)}")
     if launches:
         launch = launches[-1]
-        launch.records = records # THTODO: rm
-        launch.tensors = tensors # THTODO: rm
+        launch.records = records  # THTODO: rm
+        launch.tensors = tensors  # THTODO: rm
         print(f"Number of records: {len(launch.records)}")
         for i, record in enumerate(launch.records):
             print(f"Record {i}: {type(record).__name__}")
@@ -137,10 +137,10 @@ def matmul_kernel(lhs, rhs):
 
         traceback.print_exc()
 else:
-    print('Executing NKI JIT-ed matmul_kernel...')
-    matmul_kernel = nki.jit(matmul_kernel)
-    z2 = nki.simulate_kernel(matmul_kernel[kernel_grid], *kernel_args)
+    print("Executing NKI JIT-ed matmul_kernel...")
+    compiled_kernel = nki.jit(matmul_kernel)
+    z2 = nki.simulate_kernel(compiled_kernel[kernel_grid], *kernel_args)
 
 z1 = lhs_small @ rhs_small
-print(np.max(np.abs(z1-z2)))
+print(np.max(np.abs(z1 - z2)))
 assert np.allclose(z1, z2)
diff --git a/tests/test_nki.py b/tests/test_nki.py
index 0ec9973d..6c866ed0 100644
--- a/tests/test_nki.py
+++ b/tests/test_nki.py
@@ -2,60 +2,58 @@
 """
 Test script to verify NDArray slicing functionality after fixes
 """
-import sys
-import os
-sys.path.insert(0, os.path.join(os.path.dirname(__file__), 'triton_viz', 'core'))
-
 import numpy as np
-
-# Import the nki module directly
 from triton_viz.core.nki import NDArray
 
+# sys.path.insert(0, os.path.join(os.path.dirname(__file__), "triton_viz", "core"))
+
+
 def test_ndarray_creation():
     print("Testing NDArray creation...")
-    
+
     # Test creation with value
     data = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
-    nd1 = NDArray(value=data, name='test_array')
+    nd1 = NDArray(value=data, name="test_array")
     print(f"Created NDArray: {nd1}")
     print(f"Shape: {nd1.shape}")
     print(f"Dtype: {nd1.dtype}")
-    print(f"Value:\n{nd1.value}")
+    print(f"Value:\n{nd1.data}")
     print()
-    
+
     # Test creation with shape and dtype
-    nd2 = NDArray(shape=(2, 3), dtype=np.float32, name='shaped_array')
+    nd2 = NDArray(shape=(2, 3), dtype=np.float32, name="shaped_array")
     print(f"Created shaped NDArray: {nd2}")
     print(f"Shape: {nd2.shape}")
     print(f"Dtype: {nd2.dtype}")
     print()
 
+
 def test_slicing():
     print("Testing slicing operations...")
-    
+
     data = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
-    nd_array = NDArray(value=data, name='test_array')
-    
+    nd_array = NDArray(value=data, name="test_array")
+
     # Test [:, :] (all elements)
-    #slice_all = nd_array[:, :]
+    # slice_all = nd_array[:, :]
     slice_all = nd_array[:2, :2]
-    assert np.allclose(slice_all._value, data[:2, :2])
+    assert np.allclose(slice_all.data, data[:2, :2])
     print(f"nd_array[:, :] = {slice_all}")
-    print(f"Value:\n{slice_all.value}")
+    print(f"Value:\n{slice_all.data}")
     print()
-    
+
     # Test [:, 0] (first column)
     slice_col = nd_array[:, 0]
-    assert np.allclose(slice_col._value, data[:, 0])
+    assert np.allclose(slice_col.data, data[:, 0])
     print(f"nd_array[:, 0] = {slice_col}")
-    print(f"Value: {slice_col.value}")
+    print(f"Value: {slice_col.data}")
     print()
-    
+
     # Test [0, :] (first row)
     slice_row = nd_array[0, :]
-    assert np.allclose(slice_row._value, data[0, :])
+    assert np.allclose(slice_row.data, data[0, :])
     print(f"nd_array[0, :] = {slice_row}")
-    print(f"Value: {slice_row.value}")
+    print(f"Value: {slice_row.data}")
     print()
 
     # Test advanced indexing
@@ -64,51 +62,50 @@ def test_slicing():
         np.arange(3)[None, :],
     )
     slice_advanced = nd_array[rewritten_slice]
-    assert np.allclose(slice_advanced._value, data[rewritten_slice])
+    assert np.allclose(slice_advanced.data, data[rewritten_slice])
     print(f"nd_array[nl.arange(2)[:, None], nl.arange(3)[None, :]] = {slice_advanced}")
-    print(f"Value:\n{slice_advanced.value}")
+    print(f"Value:\n{slice_advanced.data}")
     print()
 
     # test nd_array[0, :3, :, 2:4, 2]
     data = np.reshape(np.arange(3 * 4 * 5 * 6 * 7), (3, 4, 5, 6, 7))
-    nd_array = NDArray(value=data, name='test_array_5d')
+    nd_array = NDArray(value=data, name="test_array_5d")
     rewritten_slice = (
         0,
         np.arange(3)[:, None, None],
         np.arange(5)[None, :, None],
         np.arange(2, 4)[None, None, :],
-        2
-    ) # this is what the above slice would be represented as after triton-viz tracing
+        2,
+    )  # this is what the above slice would be represented as after triton-viz tracing
     slice_advanced = nd_array[rewritten_slice]
-    assert np.allclose(
-        slice_advanced._value,
-        data[rewritten_slice]
-    )
-    assert np.allclose(slice_advanced._value, data[0, :3, :, 2:4, 2])
-    print(f"Value:\n{slice_advanced.value}")
+    assert np.allclose(slice_advanced.data, data[rewritten_slice])
+    assert np.allclose(slice_advanced.data, data[0, :3, :, 2:4, 2])
+    print(f"Value:\n{slice_advanced.data}")
     print()
 
+
 def test_arithmetic():
     print("Testing arithmetic operations...")
-    
+
     data1 = np.array([[1, 2], [3, 4]])
     data2 = np.array([[5, 6], [7, 8]])
-    
-    nd1 = NDArray(value=data1, name='array1')
-    nd2 = NDArray(value=data2, name='array2')
-    
+
+    nd1 = NDArray(value=data1, name="array1")
+    nd2 = NDArray(value=data2, name="array2")
+
     # Test addition
     result = nd1 + nd2
     print(f"Addition result: {result}")
-    print(f"Value:\n{result.value}")
+    print(f"Value:\n{result.data}")
     print()
-    
+
     # Test slicing on result
     slice_result = result[:, 0]
     print(f"Slice of result [:, 0]: {slice_result}")
-    print(f"Value: {slice_result.value}")
+    print(f"Value: {slice_result.data}")
     print()
 
+
 if __name__ == "__main__":
     test_ndarray_creation()
     test_slicing()
diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index cb837dad..048b802d 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -2,14 +2,14 @@
 
 import neuronxcc.nki.language as nl
 import inspect
-import ast
-from .nki_extract_slice import StoreCallTransformer, transform_code
+from .nki_extract_slice import transform_code
 from .nki_masked_load import masked_load, masked_store
 
 
 # Q1: slicing semantic is weird
 # Q2: why cannot we execute the .func function?
 
+
 # Multi-dimensional slice class
 class NLSlice:
     def __init__(self, start=None, stop=None, step: int = 1):
@@ -26,7 +26,7 @@ def __repr__(self):
                 stop = "None"
             if step is None:
                 step = "None"
-            repr += f"(start={start}, stop={stop}, step={step}) " 
+            repr += f"(start={start}, stop={stop}, step={step}) "
         return repr
 
     def __add__(self, other):
@@ -45,7 +45,9 @@ def __add__(self, other):
                 new_stop.append(stop + other if stop is not None else None)
                 new_step.append(step)
             return NLSlice(start=new_start, stop=new_stop, step=new_step)
-        raise TypeError(f"Unsupported operand type(s) for +: 'NLSlice' and '{type(other).__name__}'")
+        raise TypeError(
+            f"Unsupported operand type(s) for +: 'NLSlice' and '{type(other).__name__}'"
+        )
 
     def __radd__(self, other):
         new_start = []
@@ -63,7 +65,9 @@ def __radd__(self, other):
                 new_stop.append(other + stop if stop is not None else None)
                 new_step.append(step)
             return NLSlice(start=new_start, stop=new_stop, step=new_step)
-        raise TypeError(f"Unsupported operand type(s) for +: '{type(other).__name__}' and 'NLSlice'")
+        raise TypeError(
+            f"Unsupported operand type(s) for +: '{type(other).__name__}' and 'NLSlice'"
+        )
 
     def __getitem__(self, keys):
         new_start = []
@@ -77,7 +81,9 @@ def __getitem__(self, keys):
                 new_stop.append(None)
                 new_step.append(None)
             elif isinstance(k, slice):
-                assert k.start is None and k.stop is None and k.step is None, "Slice must be complete"
+                assert (
+                    k.start is None and k.stop is None and k.step is None
+                ), "Slice must be complete"
                 new_start.append(self.start[idx])
                 new_stop.append(self.stop[idx])
                 new_step.append(self.step[idx])
@@ -87,51 +93,40 @@ def __getitem__(self, keys):
 
         return NLSlice(start=new_start, stop=new_stop, step=new_step)
 
+
 class NDArray:
     def __init__(self, buffer=None, name="", **kwargs):
         self.buffer = buffer
         self.name = name
         self.kwargs = kwargs
-        value = None
+        val = None
         if "shape" in kwargs and "dtype" in kwargs:
             shape = kwargs.pop("shape")
             dtype = kwargs.pop("dtype")
-            value = np.ndarray(shape, dtype=dtype)
+            val = np.ndarray(shape, dtype=dtype)
         if "value" in kwargs:
-            assert value is None or value.shape == kwargs["value"].shape
-            value = kwargs["value"]
-        self._value = value
+            assert val is None or val.shape == kwargs["value"].shape
+            val = kwargs["value"]
+        self.data = val
 
     @property
     def shape(self):
-        return self._value.shape if self._value is not None else None
+        return self.data.shape if self.data is not None else None
 
     @property
     def dtype(self):
-        return self._value.dtype if self._value is not None else None
-
-    @property
-    def value(self):
-        return self._value
-    
-    @property
-    def data(self):
-        return self._value
-
-    @value.setter
-    def value(self, new_value):
-        self._value = new_value
+        return self.data.dtype if self.data is not None else None
 
     def data_ptr(self):
-        return self._value.ctypes.data
+        return self.data.ctypes.data
 
     def stride(self):
-        return self._value.strides
-    
+        return self.data.strides
+
     def element_size(self):
         return self.dtype.itemsize
 
-    def cpu(self): # THTODO: rm?
+    def cpu(self):  # THTODO: rm?
         return self
 
     def get_offsets(self):
@@ -144,15 +139,17 @@ def get_offsets(self):
         Returns:
             Tuple of offset arrays: (arange(A)[:, None, ..., None]*a, arange(B)[None, :, ..., None]*b, ...)
         """
-        strides = self._value.strides
-        if self._value is None:
+        strides = self.data.strides
+        if self.data is None:
             raise AttributeError("NDArray has no value - cannot compute offsets")
 
         shape = self.shape
         if len(shape) != len(strides):
-            raise ValueError(f"Shape has {len(shape)} dimensions but strides has {len(strides)} dimensions")
+            raise ValueError(
+                f"Shape has {len(shape)} dimensions but strides has {len(strides)} dimensions"
+            )
 
-        #offsets = []
+        # offsets = []
         offsets = 0
         ndim = len(shape)
 
@@ -166,10 +163,10 @@ def get_offsets(self):
 
             # Reshape and multiply by stride
             offset_array = (arange_vals * stride).reshape(broadcast_shape)
-            #offsets.append(NDArray(value=offset_array, name=f"{self.name}_offset_dim{i}"))
+            # offsets.append(NDArray(value=offset_array, name=f"{self.name}_offset_dim{i}"))
             offsets += NDArray(value=offset_array, name=f"{self.name}_offset_dim{i}")
 
-        #return tuple(offsets)
+        # return tuple(offsets)
         return offsets
 
     def __repr__(self):
@@ -177,7 +174,7 @@ def __repr__(self):
 
     def __getitem__(self, keys):
         """Implement slicing operations for NDArray"""
-        if self._value is None:
+        if self.data is None:
             raise AttributeError("NDArray has no value to slice")
         if not isinstance(keys, tuple):
             keys = (keys,)
@@ -187,36 +184,49 @@ def __getitem__(self, keys):
         arr_dim = 0
         for k in keys:
             if isinstance(k, NDArray):
-                dim_len = self._value.shape[arr_dim]
-                new_keys.append(k._value.clip(0, dim_len - 1))
+                dim_len = self.data.shape[arr_dim]
+                new_keys.append(k.data.clip(0, dim_len - 1))
             elif isinstance(k, NLSlice):
                 new_keys.append(slice(k.start, k.stop, k.step))
             elif k is None:
                 new_keys.append(k)
-                arr_dim -= 1 # add new dim -> revisit arr_dim for next key
+                arr_dim -= 1  # add new dim -> revisit arr_dim for next key
             else:
                 new_keys.append(k)
             arr_dim += 1
-        
-        sliced_value = self._value[tuple(new_keys)]
+
+        sliced_value = self.data[tuple(new_keys)]
 
         # Create a new NDArray with the sliced data
         return NDArray(value=sliced_value, name=f"{self.name}_slice")
 
-
     def _binary_op(self, other, op_func, op_name, op_symbol):
         if isinstance(other, NDArray):
-            return NDArray(value=op_func(self._value, other._value), name=f"{self.name}_{op_name}_{other.name}")
+            return NDArray(
+                value=op_func(self.data, other.data),
+                name=f"{self.name}_{op_name}_{other.name}",
+            )
         elif np.isscalar(other):
-            return NDArray(value=op_func(self._value, other), name=f"{self.name}_{op_name}_scalar")
-        raise TypeError(f"Unsupported operand type(s) for {op_symbol}: 'NDArray' and '{type(other).__name__}'")
+            return NDArray(
+                value=op_func(self.data, other), name=f"{self.name}_{op_name}_scalar"
+            )
+        raise TypeError(
+            f"Unsupported operand type(s) for {op_symbol}: 'NDArray' and '{type(other).__name__}'"
+        )
 
     def _rbinary_op(self, other, op_func, op_name, op_symbol):
         if isinstance(other, NDArray):
-            return NDArray(value=op_func(other._value, self._value), name=f"{other.name}_{op_name}_{self.name}")
+            return NDArray(
+                value=op_func(other.data, self.data),
+                name=f"{other.name}_{op_name}_{self.name}",
+            )
         elif np.isscalar(other):
-            return NDArray(value=op_func(other, self._value), name=f"scalar_{op_name}_{self.name}")
-        raise TypeError(f"Unsupported operand type(s) for {op_symbol}: '{type(other).__name__}' and 'NDArray'")
+            return NDArray(
+                value=op_func(other, self.data), name=f"scalar_{op_name}_{self.name}"
+            )
+        raise TypeError(
+            f"Unsupported operand type(s) for {op_symbol}: '{type(other).__name__}' and 'NDArray'"
+        )
 
     # Define operator +/-/*//
     def __add__(self, other):
@@ -294,12 +304,14 @@ def ndarray(self, shape, dtype, *, buffer=None, name=None, **kwargs):
                 ret = self.shared_hbm_arrays[name]
             else:
                 # Create a new shared HBM array and store it
-                ret = NDArray(buffer=buffer, name=name, shape=shape, dtype=dtype, **kwargs)
+                ret = NDArray(
+                    buffer=buffer, name=name, shape=shape, dtype=dtype, **kwargs
+                )
                 self.shared_hbm_arrays[name] = ret
         else:
             ret = NDArray(buffer=buffer, name=name, shape=shape, dtype=dtype, **kwargs)
         return ret
-    
+
     def zeros(self, shape, dtype, *, buffer=None, name=None, **kwargs):
         value = np.zeros(shape, dtype=dtype)
         return self.ndarray(
@@ -307,14 +319,6 @@ def zeros(self, shape, dtype, *, buffer=None, name=None, **kwargs):
         )
 
     def arange(self, *args):
-        if len(args) == 1:
-            start = 0
-            stop = args[0]
-        elif len(args) == 2:
-            start = args[0]
-            stop = args[1]
-        else:
-            raise ValueError("arange expects 1 or 2 arguments")
         return NDArray(value=np.arange(*args))
 
     def program_id(self, axis: int):
@@ -328,15 +332,15 @@ def program_id(self, axis: int):
             raise ValueError(f"Invalid axis: {axis}. Must be 0, 1, or 2.")
 
     def load(self, src: NDArray, *, mask=None, dtype=None, **kwargs):
-        value = src._value
+        value = src.data
         if isinstance(mask, NDArray):
-            value = value[mask._value]
+            value = value[mask.data]
         elif mask is not None:
             value = value[mask]
         if dtype is not None:
             value = value.astype(dtype)
 
-        mask_value = getattr(mask, "_value", np.ones_like(src))
+        mask_value = getattr(mask, "data", np.ones_like(src))
         new_shape = []
         for i, _v in enumerate(mask_value.shape):
             assert np.unique(mask_value.sum(i)).size <= 2
@@ -348,11 +352,13 @@ def load(self, src: NDArray, *, mask=None, dtype=None, **kwargs):
 
     def load_transpose2d(self, src: NDArray, *, mask=None, dtype=None, **kwargs):
         # THTODO
-        value = src._value
-        return self.load(NDArray(value=value.T, name=src.name), mask=mask, dtype=dtype, **kwargs)
+        value = src.data
+        return self.load(
+            NDArray(value=value.T, name=src.name), mask=mask, dtype=dtype, **kwargs
+        )
 
     def store(self, dst: NDArray, value: NDArray, *, mask=None, **kwargs):
-        dst._value[mask._value] = value._value.ravel()
+        dst.data[mask.data] = value.data.ravel()
         return dst
 
     def _convert_keys_to_numpy(self, keys):
@@ -360,15 +366,15 @@ def _convert_keys_to_numpy(self, keys):
         if isinstance(keys, (tuple, list)):
             return tuple(self._convert_keys_to_numpy(k) for k in keys)
         elif isinstance(keys, NDArray):
-            return keys._value
+            return keys.data
         else:
             return keys
 
     def masked_load(self, src: NDArray, keys, *, mask=None, **kwargs):
         """Load array elements with masking for out-of-bounds errors."""
         # Convert NDArray to numpy array
-        ndarray = src._value
-        mask_value = getattr(mask, "_value", mask) if mask is not None else None
+        ndarray = src.data
+        mask_value = getattr(mask, "data", mask) if mask is not None else None
 
         # Convert any NDArrays in keys to numpy arrays
         numpy_keys = self._convert_keys_to_numpy(keys)
@@ -382,9 +388,9 @@ def masked_load(self, src: NDArray, keys, *, mask=None, **kwargs):
     def masked_store(self, dst: NDArray, keys, value: NDArray, *, mask=None, **kwargs):
         """Store array elements with masking for out-of-bounds errors."""
         # Convert NDArrays to numpy arrays
-        ndarray = dst._value
-        value_array = value._value
-        mask_value = getattr(mask, "_value", mask) if mask is not None else None
+        ndarray = dst.data
+        value_array = value.data
+        mask_value = getattr(mask, "data", mask) if mask is not None else None
 
         # Convert any NDArrays in keys to numpy arrays
         numpy_keys = self._convert_keys_to_numpy(keys)
@@ -394,8 +400,8 @@ def masked_store(self, dst: NDArray, keys, value: NDArray, *, mask=None, **kwarg
 
         return dst
 
-    def unary_op(self, x: NDArray, np_func, op_name, **kwargs):
-        return NDArray(value=np_func(x._value), name=f"{x.name}_{op_name}", **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)
 
     # Elementwise operator implementations
     def exp(self, x: NDArray, **kwargs):
@@ -412,14 +418,16 @@ def tanh(self, x: NDArray, **kwargs):
 
     def silu(self, x: NDArray, **kwargs):
         # SiLU(x) = x * sigmoid(x)
-        sigmoid_x = 1 / (1 + np.exp(-x._value))
-        return NDArray(value=x._value * sigmoid_x, name=f"{x.name}_silu", **kwargs)
+        sigmoid_x = 1 / (1 + np.exp(-x.data))
+        return NDArray(value=x.data * sigmoid_x, name=f"{x.name}_silu", **kwargs)
 
     def gelu(self, x: NDArray, **kwargs):
         # GELU(x) = 0.5 * x * (1 + tanh(sqrt(2/pi) * (x + 0.044715 * x^3)))
         sqrt_2_pi = np.sqrt(2 / np.pi)
-        inner = sqrt_2_pi * (x._value + 0.044715 * np.power(x._value, 3))
-        return NDArray(value=0.5 * x._value * (1 + np.tanh(inner)), name=f"{x.name}_gelu", **kwargs)
+        inner = sqrt_2_pi * (x.data + 0.044715 * np.power(x.data, 3))
+        return NDArray(
+            value=0.5 * x.data * (1 + np.tanh(inner)), name=f"{x.name}_gelu", **kwargs
+        )
 
     def sqrt(self, x: NDArray, **kwargs):
         return self._unary_op(x, np.sqrt, "sqrt", **kwargs)
@@ -432,9 +440,17 @@ def log(self, x: NDArray, **kwargs):
 
     def pow(self, x: NDArray, exponent, **kwargs):
         if isinstance(exponent, NDArray):
-            return NDArray(value=np.power(x._value, exponent._value), name=f"{x.name}_pow_{exponent.name}", **kwargs)
+            return NDArray(
+                value=np.power(x.data, exponent.data),
+                name=f"{x.name}_pow_{exponent.name}",
+                **kwargs,
+            )
         elif np.isscalar(exponent):
-            return NDArray(value=np.power(x._value, exponent), name=f"{x.name}_pow_{exponent}", **kwargs)
+            return NDArray(
+                value=np.power(x.data, exponent),
+                name=f"{x.name}_pow_{exponent}",
+                **kwargs,
+            )
         else:
             raise TypeError(f"Unsupported exponent type: {type(exponent)}")
 
@@ -442,18 +458,21 @@ def reciprocal(self, x: NDArray, **kwargs):
         return self._unary_op(x, lambda v: 1 / v, "reciprocal", **kwargs)
 
     def matmul(self, x: NDArray, y: NDArray, transpose_x=False, mask=None, **kwargs):
-        x_value = x._value 
+        x_value = x.data
         if transpose_x:
             x_value = x_value.T
-        y_value = y._value 
-        return NDArray(value=(x_value @ y_value), name=f"{x.name}_{y.name}_matmul", **kwargs)
+        y_value = y.data
+        return NDArray(
+            value=(x_value @ y_value), name=f"{x.name}_{y.name}_matmul", **kwargs
+        )
 
     def copy(self, x: NDArray, **kwargs):
-        return self.unary_op(x, np.copy, "copy", **kwargs)
+        return self._unary_op(x, np.copy, "copy", **kwargs)
 
     def range(self, stop):
         return range(stop)
 
+
 nki_builder = Builder()
 
 
@@ -476,17 +495,17 @@ def patch():
     nl.affine_range = nki_builder.range
     nl.par_dim
     nl.zeros = nki_builder.zeros
-    nl.mgrid = NDArray(value=np.mgrid, buffer=nl.sbuf, name='mgrid')
+    nl.mgrid = NDArray(value=np.mgrid, buffer=nl.sbuf, name="mgrid")
     nl.matmul = nki_builder.matmul
     nl.copy = nki_builder.copy
 
     # attention-specific
     nl.load_transpose2d = nki_builder.load_transpose2d
-    #nisa.affine_select
-    #nl.tensor_reduce
-    #nisa.activation
+    # nisa.affine_select
+    # nl.tensor_reduce
+    # nisa.activation
     nl.broadcast_to
-    #nisa.nc_transpose
+    # nisa.nc_transpose
 
     # Elementwise operators
     nl.exp = nki_builder.exp
@@ -516,52 +535,54 @@ def __init__(self, fn):
         self.fn = fn
 
     def run(self, *args, **kwargs):
-        grid_dims = kwargs.pop("grid", (1, 1, 1))  # Remove grid from kwargs to avoid passing it to the function
+        grid_dims = kwargs.pop(
+            "grid", (1, 1, 1)
+        )  # Remove grid from kwargs to avoid passing it to the function
         # make it 3d if not
         if len(grid_dims) == 1:
             grid_dims = (grid_dims[0], 1, 1)
         elif len(grid_dims) == 2:
             grid_dims = (grid_dims[0], grid_dims[1], 1)
         elif len(grid_dims) != 3:
-            raise ValueError(f"Grid must be 1, 2, or 3 dimensions, got {len(grid_dims)}")
+            raise ValueError(
+                f"Grid must be 1, 2, or 3 dimensions, got {len(grid_dims)}"
+            )
         nki_builder.set_grid_dim(grid_dims)
         nki_builder.shared_hbm_arrays = {}
         nki_builder.fn = self.fn
 
         kwargs.pop("warmup", None)  # Remove warmup from kwargs if it exists
-        client_manager = kwargs.pop("client_manager", None)  # Remove client_manager from kwargs if it exists
+        client_manager = kwargs.pop(
+            "client_manager", None
+        )  # Remove client_manager from kwargs if it exists
 
         # Call grid_callback once before grid execution (similar to Triton)
         if client_manager is not None:
             client_manager.grid_callback(grid_dims)
 
-        patch() # NKI interpreter patching
+        patch()  # NKI interpreter patching
 
-        #with client_manager.patch():
-        #v
-        #kwargs.update({"client_manager": client_manager})
-        #ret = self.interpreter_fn.run(*args, **kwargs)
-        #self.finalize()
-        #return ret
+        # with client_manager.patch():
+        # v
 
         # Apply AST transformer to convert nl.load/nl.store calls to nl.masked_load/nl.masked_store
-        import types
-        if hasattr(self.fn, '__code__'):
+        if hasattr(self.fn, "__code__"):
             # Get the source code of the function
             source_code = inspect.getsource(self.fn)
             # Transform the source code using the AST transformer
             transformed_code = transform_code(source_code)
             # Create a new function from the transformed code
-            #exec_globals = self.fn.__globals__.copy()
-            #exec(transformed_code, exec_globals)
-            ## Replace the original function with the transformed one
-            #self.fn = exec_globals[self.fn.__name__]
             exec_globals = self.fn.__globals__.copy()
-            import random, string, os
-            rand_str = ''.join(random.choices(string.ascii_letters + string.digits, k=16))
-            os.makedirs('/tmp/triton-viz', exist_ok=True)
-            filename = f'/tmp/triton-viz/{rand_str}.py'
-            with open(filename, 'w') as f:
+            import random
+            import string
+            import os
+
+            rand_str = "".join(
+                random.choices(string.ascii_letters + string.digits, k=16)
+            )
+            os.makedirs("/tmp/triton-viz", exist_ok=True)
+            filename = f"/tmp/triton-viz/{rand_str}.py"
+            with open(filename, "w") as f:
                 f.write(transformed_code)
             code_obj = compile(transformed_code, filename=filename, mode="exec")
             exec(code_obj, exec_globals)
@@ -579,7 +600,7 @@ def run(self, *args, **kwargs):
                         client_manager.grid_idx_callback((x, y, z))
 
                     result = self.fn(*args, **kwargs)
-        #^
+        # ^
 
         unpatch()
-        return result.value
+        return result.data

From 0fb3dbc5adc481edf7382973e48f1ee447bde3aa Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Sat, 25 Oct 2025 09:13:51 -0400
Subject: [PATCH 54/90] make nki tracing work again

---
 triton_viz/core/trace.py | 37 +++++++++++++++++++++++++++++--------
 1 file changed, 29 insertions(+), 8 deletions(-)

diff --git a/triton_viz/core/trace.py b/triton_viz/core/trace.py
index 5c8ce7e2..ac6a4467 100644
--- a/triton_viz/core/trace.py
+++ b/triton_viz/core/trace.py
@@ -19,7 +19,8 @@ def dummy_benchmarker(fn, quantiles):
     fn()
     return (1.0, 1.0, 1.0)
 
-class TraceInterface():
+
+class TraceInterface:
     def __init__(self, client: Union[str, Client]) -> None:
         self.client_manager = ClientManager()
         self.add_client(client)
@@ -47,6 +48,7 @@ def finalize(self):
         self.client_manager.finalize()
         launches.append(self.client_manager.launch)
 
+
 class TritonTrace(KernelInterface, TraceInterface):
     def __init__(
         self,
@@ -70,7 +72,7 @@ def unpack_kernel(
             if isinstance(source, InterpretedFunction):
                 return None, source.fn, source
             raise TypeError(f"Unsupported runner type: {type(source)}")
-        
+
         if isinstance(runner, Autotuner):
             self.jit_fn, self.base_fn, self.interpreted_fn = unpack_kernel(runner.fn)
             # replace the benchmark with a dummy that just calls the function once
@@ -86,7 +88,7 @@ def unpack_kernel(
             self.jit_fn, self.base_fn, self.interpreted_fn = unpack_kernel(runner)
             self.runner = self.interpreted_fn
             self.warmup_runner = self.jit_fn
-        
+
         self.arg_names = runner.arg_names
 
         self.fn = runner
@@ -97,7 +99,7 @@ def warmup(self, *args, **kwargs):
         with self.client_manager.patch_warmup(self.jit_fn):
             if self.warmup_runner:
                 self.warmup_runner.warmup(*args, **kwargs)
-    
+
     def run(self, *args, **kwargs):
         with self.client_manager.patch_warmup(self.jit_fn):
             if self.warmup_runner:
@@ -109,7 +111,7 @@ def run(self, *args, **kwargs):
             ret = self.runner.run(*args, **kwargs)
             self.finalize()
             return ret
-    
+
     def __call__(self, *args, **kwargs):
         # When a traced JIT function is called from within another JIT function,
         # we need to execute the underlying function directly
@@ -129,6 +131,7 @@ def __init__(self, kernel, client: str | Client) -> None:
             self.func = kernel.func
         elif isinstance(kernel, NKIInterpretedFunction):
             self.interpreter_fn = kernel
+            self.func = kernel.fn
         else:
             self.interpreter_fn = NKIInterpretedFunction(kernel)
             self.func = kernel
@@ -137,10 +140,16 @@ def __init__(self, kernel, client: str | Client) -> None:
 
     def __getitem__(self, *grid):
         return KernelInterface.__getitem__(self, tuple(*grid))
-    
+
     def __call__(self, *args, **kwargs):
         return self[(1, 1, 1)](*args, **kwargs)
 
+    def run(self, *args, **kwargs):
+        kwargs.update({"client_manager": self.client_manager})
+        ret = self.interpreter_fn.run(*args, **kwargs)
+        self.finalize()
+        return ret
+
 
 def trace(clients: Union[str, Client, None] = None, backend: str = "triton"):
     """
@@ -169,14 +178,26 @@ def decorator(kernel) -> TraceInterface:
             else:
                 raise ValueError(f"Unknown backend: {backend}")
 
+        # Handle NKI functions specifically
+        if backend == "nki":
+            from .nki import NKIInterpretedFunction
+
+            if isinstance(kernel, NKIInterpretedFunction):
+                return NKITrace(kernel, clients)
+            else:
+                # Wrap plain functions as NKIInterpretedFunction
+                interpreted_fn = NKIInterpretedFunction(kernel)
+                return NKITrace(interpreted_fn, clients)
+
         # If the object is already initialized as a TraceInterface, just append the new client(s)
         if isinstance(kernel, TraceInterface):
             trace = kernel
             trace.add_client(clients)
             return trace
-        
 
-        raise TypeError(f"Expected JITFunction, InterpretedFunction or Trace, got {type(kernel)}")
+        raise TypeError(
+            f"Expected JITFunction, InterpretedFunction or Trace, got {type(kernel)}"
+        )
 
     return decorator
 

From a225146b127ad01a8d9e98f2bc962847b716bc55 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Sun, 26 Oct 2025 19:16:42 -0400
Subject: [PATCH 55/90] don't immediately exit web server when share=false

---
 triton_viz/visualizer/interface.py | 10 +++-------
 1 file changed, 3 insertions(+), 7 deletions(-)

diff --git a/triton_viz/visualizer/interface.py b/triton_viz/visualizer/interface.py
index 9fa30b5d..c815cd22 100644
--- a/triton_viz/visualizer/interface.py
+++ b/triton_viz/visualizer/interface.py
@@ -330,14 +330,10 @@ def launch(share: bool = True, port: int | None = None):
         print("--------")
         _server_state.set_local_port(actual_port)
 
-        # Run Flask in a background thread so callers can continue (non-blocking)
-        def _run_local():
-            app.run(host="0.0.0.0", port=actual_port, debug=True, use_reloader=False)
+        # For share=False, we want to block and keep the server running
+        # This is the traditional behavior for local web servers
+        app.run(host="0.0.0.0", port=actual_port, debug=True, use_reloader=False)
 
-        flask_thread = threading.Thread(target=_run_local, daemon=True)
-        flask_thread.start()
-        # Give the server a moment to bind the port
-        time.sleep(0.5)
         return local_url, None
 
 

From 6571b71f71ab6ab5ad10e62d3cb57b36e22b9cbf Mon Sep 17 00:00:00 2001
From: gujialiang123 <562960306@qq.com>
Date: Sun, 26 Oct 2025 19:41:00 -0400
Subject: [PATCH 56/90] feat(flow): embed memory-flow badges in Load/Store/Dot;
 add Flow Diagram view; fix NKI matmul Dot shapes; robust value APIs; 2D shape
 mapping; add NKI record metadata

---
 nki-examples/matmul.py             | 141 ++++++++++++++++++++---------
 triton_viz/core/nki_masked_load.py |  43 ++++++---
 triton_viz/core/trace.py           |  38 +++++---
 triton_viz/static/gridblock.js     |  20 ++++
 triton_viz/static/load.js          |  21 +++++
 triton_viz/static/load_utils.js    |  33 ++++++-
 triton_viz/static/matmul.js        |  74 +++++++++++----
 triton_viz/static/nki.js           | 107 ++++++++++++++++++++++
 triton_viz/static/store.js         |  21 +++++
 triton_viz/visualizer/draw.py      |  51 ++++++++++-
 triton_viz/visualizer/interface.py |  73 ++++++++++-----
 11 files changed, 507 insertions(+), 115 deletions(-)
 create mode 100644 triton_viz/static/nki.js

diff --git a/nki-examples/matmul.py b/nki-examples/matmul.py
index 4e3c458b..c11cd54d 100644
--- a/nki-examples/matmul.py
+++ b/nki-examples/matmul.py
@@ -1,18 +1,14 @@
 from neuronxcc import nki
 import neuronxcc.nki.language as nl
-
-import tp
-import torch
 import triton_viz
-import torch_xla.core.xla_model as xm
 from triton_viz.clients import Tracer
-from triton_viz.core import config as cfg
 from triton_viz.core.trace import launches
 import numpy as np
-import math
 
-tensors = [] # THTODO: rm
-records = [triton_viz.core.data.Grid((0,0,0))] # THTODO: rm
+tensors = []  # THTODO: rm
+records = [triton_viz.core.data.Grid((0, 0, 0))]  # THTODO: rm
+
+
 def matmul_kernel(lhs, rhs):
     """NKI matmul_kernel to compute a matrix multiplication operation in a tiled manner
 
@@ -40,8 +36,7 @@ def matmul_kernel(lhs, rhs):
     tensors.append(rhs)
     tensors.append(result)
     # Use affine_range to loop over tiles
-    m_range = nl.arange(M)[:, None]
-    n_range = nl.arange(N)[None, :]
+    time_idx = 0
     for m in nl.affine_range(M // TILE_M):
         for n in nl.affine_range(N // TILE_N):
             # Allocate a tensor in PSUM
@@ -57,33 +52,80 @@ def matmul_kernel(lhs, rhs):
                 lhs_f = nl.arange(TILE_K)[None, :] + k * TILE_K
                 lhs_mask = (lhs_p < M) & (lhs_f < K)
                 lhs_tile = nl.load(lhs[lhs_p, lhs_f], mask=lhs_mask)
-                records.append(triton_viz.core.data.Load(
-                    lhs.data_ptr(),
-                    (4 * (lhs_p * K + lhs_f))._value,
-                    lhs_mask._value
-                )) # THTODO: rm
+                records.append(
+                    triton_viz.core.data.Load(
+                        lhs.data_ptr(),
+                        (4 * (lhs_p * K + lhs_f))._value,
+                        lhs_mask._value,
+                    )
+                )  # THTODO: rm
+                # NKI meta for viz (HBM->SBUF)
+                try:
+                    bytes_lhs = int(np.count_nonzero(lhs_mask._value)) * int(
+                        lhs.dtype.itemsize
+                    )
+                except Exception:
+                    bytes_lhs = 0
+                records[-1].mem_src = "HBM"
+                records[-1].mem_dst = "SBUF"
+                records[-1].bytes = bytes_lhs
+                records[-1].time_idx = time_idx
+                records[-1].backend = "nki"
+                time_idx += 1
 
                 rhs_p = nl.arange(TILE_K)[:, None] + k * TILE_K
                 rhs_f = nl.arange(TILE_N)[None, :] + n * TILE_N
-                rhs_mask = (rhs_p < K) & (rhs_f < N) 
+                rhs_mask = (rhs_p < K) & (rhs_f < N)
                 rhs_tile = nl.load(rhs[rhs_p, rhs_f], mask=rhs_mask)
-                records.append(triton_viz.core.data.Load(
-                    rhs.data_ptr(),
-                    (4 * (rhs_p * N + rhs_f))._value,
-                    rhs_mask._value
-                )) # THTODO: rm
+                records.append(
+                    triton_viz.core.data.Load(
+                        rhs.data_ptr(),
+                        (4 * (rhs_p * N + rhs_f))._value,
+                        rhs_mask._value,
+                    )
+                )  # THTODO: rm
+                try:
+                    bytes_rhs = int(np.count_nonzero(rhs_mask._value)) * int(
+                        rhs.dtype.itemsize
+                    )
+                except Exception:
+                    bytes_rhs = 0
+                records[-1].mem_src = "HBM"
+                records[-1].mem_dst = "SBUF"
+                records[-1].bytes = bytes_rhs
+                records[-1].time_idx = time_idx
+                records[-1].backend = "nki"
+                time_idx += 1
 
                 # Accumulate partial-sums into PSUM
                 x = nl.matmul(lhs_tile[...], rhs_tile[...], transpose_x=False)
                 res_psum += x
-                records.append(triton_viz.core.data.Dot(
-                    lhs_tile.shape,
-                    rhs_tile.shape,
-                    res_psum.shape,
-                    lhs_tile._value,
-                    rhs_tile._value,
-                    x._value
-                )) # THTODO: rm
+
+                # For visualization, use standard matmul convention A(M,K)·B(K,N)->C(M,N)
+                # Our lhs_tile is (K, M) physically; transpose to (M, K) for Dot record
+                lhs_viz = lhs_tile._value.T
+                rhs_viz = rhs_tile._value
+                lhs_shape_viz = (lhs_tile.shape[1], lhs_tile.shape[0])  # (M, K)
+                rhs_shape_viz = rhs_tile.shape  # (K, N)
+                out_shape_viz = res_psum.shape  # (M, N)
+
+                records.append(
+                    triton_viz.core.data.Dot(
+                        lhs_shape_viz,
+                        rhs_shape_viz,
+                        out_shape_viz,
+                        lhs_viz,
+                        rhs_viz,
+                        x._value,
+                    )
+                )  # THTODO: rm
+                records[-1].mem_src = "SBUF"
+                records[-1].mem_dst = "PSUM"
+                records[-1].tile_shape = (TILE_M, TILE_N)
+                records[-1].k = TILE_K
+                records[-1].time_idx = time_idx
+                records[-1].backend = "nki"
+                time_idx += 1
 
             # Copy the result from PSUM back to SBUF, and cast to expected output data-type
             res_sb = nl.copy(res_psum, dtype=result.dtype)
@@ -91,16 +133,31 @@ def matmul_kernel(lhs, rhs):
             out_p = nl.arange(TILE_M)[:, None] + m * TILE_M
             out_f = nl.arange(TILE_N)[None, :] + n * TILE_N
             out_mask = (out_p < M) & (out_f < N)
-            nl.store(result[m * TILE_M:(m + 1) * TILE_M, n * TILE_N:(n + 1) * TILE_N],
-                     value=res_sb)
-            records.append(triton_viz.core.data.Store(
-                result.data_ptr(),
-                (4 * (out_p * N + out_f))._value,
-                out_mask._value
-            )) # THTODO: rm
+            nl.store(
+                result[m * TILE_M : (m + 1) * TILE_M, n * TILE_N : (n + 1) * TILE_N],
+                value=res_sb,
+            )
+            records.append(
+                triton_viz.core.data.Store(
+                    result.data_ptr(), (4 * (out_p * N + out_f))._value, out_mask._value
+                )
+            )  # THTODO: rm
+            try:
+                bytes_out = int(np.count_nonzero(out_mask._value)) * int(
+                    result.dtype.itemsize
+                )
+            except Exception:
+                bytes_out = 0
+            records[-1].mem_src = "SBUF"
+            records[-1].mem_dst = "HBM"
+            records[-1].bytes = bytes_out
+            records[-1].time_idx = time_idx
+            records[-1].backend = "nki"
+            time_idx += 1
 
     return result
 
+
 TRITON_VIZ = True
 kernel_grid = (1, 1, 1)
 lhs_small = np.arange(16).astype(np.float32).reshape(4, 4)
@@ -108,7 +165,7 @@ def matmul_kernel(lhs, rhs):
 kernel_args = (lhs_small, rhs_small)
 
 if TRITON_VIZ:
-    print('Executing matmul_kernel with NKI interpreter...')
+    print("Executing matmul_kernel with NKI interpreter...")
     matmul_kernel = triton_viz.trace(clients=Tracer(), backend="nki")(matmul_kernel)
     kk = matmul_kernel[kernel_grid]
     z2 = kk(*kernel_args)
@@ -116,8 +173,8 @@ def matmul_kernel(lhs, rhs):
     print(f"Number of launches: {len(launches)}")
     if launches:
         launch = launches[-1]
-        launch.records = records # THTODO: rm
-        launch.tensors = tensors # THTODO: rm
+        launch.records = records  # THTODO: rm
+        launch.tensors = tensors  # THTODO: rm
         print(f"Number of records: {len(launch.records)}")
         for i, record in enumerate(launch.records):
             print(f"Record {i}: {type(record).__name__}")
@@ -130,17 +187,17 @@ def matmul_kernel(lhs, rhs):
 
     # Try to launch visualization
     try:
-        triton_viz.launch(share=False)
+        triton_viz.launch(share=True, port=8005)
     except Exception as e:
         print(f"\nError during visualization: {e}")
         import traceback
 
         traceback.print_exc()
 else:
-    print('Executing NKI JIT-ed matmul_kernel...')
+    print("Executing NKI JIT-ed matmul_kernel...")
     matmul_kernel = nki.jit(matmul_kernel)
     z2 = nki.simulate_kernel(matmul_kernel[kernel_grid], *kernel_args)
 
 z1 = lhs_small @ rhs_small
-print(np.max(np.abs(z1-z2)))
+print(np.max(np.abs(z1 - z2)))
 assert np.allclose(z1, z2)
diff --git a/triton_viz/core/nki_masked_load.py b/triton_viz/core/nki_masked_load.py
index 1094756b..5faab7cd 100644
--- a/triton_viz/core/nki_masked_load.py
+++ b/triton_viz/core/nki_masked_load.py
@@ -2,8 +2,9 @@
 
 get = lambda x, default: default if x is None else x
 
+
 def normalize_slice(ndarray: np.ndarray, keys: tuple) -> tuple[tuple, tuple]:
-    """ Separate singleton dims (None) and add bounds to slices like :2, 2:, or :"""
+    """Separate singleton dims (None) and add bounds to slices like :2, 2:, or :"""
     singleton_dims = []
     dim_idx = 0
     new_keys = []
@@ -28,7 +29,10 @@ def normalize_slice(ndarray: np.ndarray, keys: tuple) -> tuple[tuple, tuple]:
         dim_idx += 1
     return tuple(new_keys), tuple(singleton_dims)
 
-def _calculate_target_shape(keys: tuple, original_shape: tuple[int, ...]) -> tuple[int, ...]:
+
+def _calculate_target_shape(
+    keys: tuple, original_shape: tuple[int, ...]
+) -> tuple[int, ...]:
     """Calculate the shape needed for the padded array to handle indexing."""
     target_shape = list(original_shape)
     for i, (arr_dim, key) in enumerate(zip(original_shape, keys)):
@@ -40,7 +44,10 @@ def _calculate_target_shape(keys: tuple, original_shape: tuple[int, ...]) -> tup
             target_shape[i] = max(arr_dim, np.max(key) + 1)
     return tuple(target_shape)
 
-def _get_valid_indices(keys: tuple, original_shape: tuple[int, ...], result_shape: tuple[int, ...]) -> np.ndarray:
+
+def _get_valid_indices(
+    keys: tuple, original_shape: tuple[int, ...], result_shape: tuple[int, ...]
+) -> np.ndarray:
     """Get a boolean mask indicating which result indices are within original array bounds."""
     coords = np.mgrid[[slice(0, s) for s in result_shape]]
     valid_mask = np.ones(result_shape, dtype=bool)
@@ -55,11 +62,14 @@ def _get_valid_indices(keys: tuple, original_shape: tuple[int, ...], result_shap
             valid_indices = np.array(key)
             valid_mask &= (0 <= valid_indices) & (valid_indices < arr_dim)
         elif isinstance(key, (int, np.integer)):
-            valid_mask &= (0 <= int(key) < arr_dim)
+            valid_mask &= 0 <= int(key) < arr_dim
 
     return valid_mask
 
-def masked_load(ndarray: np.ndarray, keys: tuple, mask: np.ndarray = None) -> np.ndarray:
+
+def masked_load(
+    ndarray: np.ndarray, keys: tuple, mask: np.ndarray = None
+) -> np.ndarray:
     """
     Load array elements with masking for out-of-bounds errors.
 
@@ -106,14 +116,19 @@ def masked_load(ndarray: np.ndarray, keys: tuple, mask: np.ndarray = None) -> np
     if len(oob_coords[0]) > 0:
         # Get the first OOB coordinate
         oob_idx = tuple(coord[0] for coord in oob_coords)
-        raise IndexError(f"index {oob_idx} is out of bounds for array of size {ndarray.shape}")
+        raise IndexError(
+            f"index {oob_idx} is out of bounds for array of size {ndarray.shape}"
+        )
 
     valid_mask = mask & in_bounds_mask
     result[~valid_mask] = 6
 
     return np.expand_dims(result, singleton_dims)
 
-def masked_store(ndarray: np.ndarray, keys: tuple, value: np.ndarray, mask: np.ndarray = None) -> None:
+
+def masked_store(
+    ndarray: np.ndarray, keys: tuple, value: np.ndarray, mask: np.ndarray = None
+) -> None:
     """
     Store array elements with masking for out-of-bounds errors.
 
@@ -130,7 +145,7 @@ def masked_store(ndarray: np.ndarray, keys: tuple, value: np.ndarray, mask: np.n
     """
     # Handle mask=None case
     if mask is None:
-        import tp; tp.log(f'{ndarray.shape=}, {keys=}')
+        # Debug logging removed: avoid external dependencies
         ndarray[keys] = value
         return
 
@@ -157,14 +172,16 @@ def masked_store(ndarray: np.ndarray, keys: tuple, value: np.ndarray, mask: np.n
     if len(oob_coords[0]) > 0:
         # Get the first OOB coordinate
         oob_idx = tuple(coord[0] for coord in oob_coords)
-        raise IndexError(f"index {oob_idx} is out of bounds for array of size {ndarray.shape}")
+        raise IndexError(
+            f"index {oob_idx} is out of bounds for array of size {ndarray.shape}"
+        )
 
-    print('!Ndarray')
+    print("!Ndarray")
     print(ndarray)
-    print('!values')
+    print("!values")
     print(value)
-    print('!Offsets')
+    print("!Offsets")
     print(offsets)
-    print('!mask')
+    print("!mask")
     print(mask)
     ndarray.ravel()[offsets[mask]] = value.ravel()[offsets[mask]]
diff --git a/triton_viz/core/trace.py b/triton_viz/core/trace.py
index 5c8ce7e2..3ef209b8 100644
--- a/triton_viz/core/trace.py
+++ b/triton_viz/core/trace.py
@@ -19,7 +19,8 @@ def dummy_benchmarker(fn, quantiles):
     fn()
     return (1.0, 1.0, 1.0)
 
-class TraceInterface():
+
+class TraceInterface:
     def __init__(self, client: Union[str, Client]) -> None:
         self.client_manager = ClientManager()
         self.add_client(client)
@@ -47,6 +48,7 @@ def finalize(self):
         self.client_manager.finalize()
         launches.append(self.client_manager.launch)
 
+
 class TritonTrace(KernelInterface, TraceInterface):
     def __init__(
         self,
@@ -70,7 +72,7 @@ def unpack_kernel(
             if isinstance(source, InterpretedFunction):
                 return None, source.fn, source
             raise TypeError(f"Unsupported runner type: {type(source)}")
-        
+
         if isinstance(runner, Autotuner):
             self.jit_fn, self.base_fn, self.interpreted_fn = unpack_kernel(runner.fn)
             # replace the benchmark with a dummy that just calls the function once
@@ -86,7 +88,7 @@ def unpack_kernel(
             self.jit_fn, self.base_fn, self.interpreted_fn = unpack_kernel(runner)
             self.runner = self.interpreted_fn
             self.warmup_runner = self.jit_fn
-        
+
         self.arg_names = runner.arg_names
 
         self.fn = runner
@@ -97,7 +99,7 @@ def warmup(self, *args, **kwargs):
         with self.client_manager.patch_warmup(self.jit_fn):
             if self.warmup_runner:
                 self.warmup_runner.warmup(*args, **kwargs)
-    
+
     def run(self, *args, **kwargs):
         with self.client_manager.patch_warmup(self.jit_fn):
             if self.warmup_runner:
@@ -109,7 +111,7 @@ def run(self, *args, **kwargs):
             ret = self.runner.run(*args, **kwargs)
             self.finalize()
             return ret
-    
+
     def __call__(self, *args, **kwargs):
         # When a traced JIT function is called from within another JIT function,
         # we need to execute the underlying function directly
@@ -135,9 +137,18 @@ def __init__(self, kernel, client: str | Client) -> None:
 
         TraceInterface.__init__(self, client)
 
-    def __getitem__(self, *grid):
-        return KernelInterface.__getitem__(self, tuple(*grid))
-    
+    def run(self, *args, **kwargs):
+        kwargs.update({"client_manager": self.client_manager})
+        ret = self.interpreter_fn.run(*args, **kwargs)
+        self.finalize()
+        return ret
+
+    def __getitem__(self, grid):
+        if isinstance(grid, tuple):
+            return KernelInterface.__getitem__(self, grid)
+        else:
+            return KernelInterface.__getitem__(self, (grid,))
+
     def __call__(self, *args, **kwargs):
         return self[(1, 1, 1)](*args, **kwargs)
 
@@ -161,11 +172,13 @@ def decorator(kernel) -> TraceInterface:
             return kernel
 
         # First-time wrapping
+        # Triton backend need JIT/Interpreter/Autotuner；
+        # NKI allow Python function（ NKIInterpretedFunction）
+        if backend == "nki":
+            return NKITrace(kernel, clients)
         if isinstance(kernel, (JITFunction, InterpretedFunction, Autotuner)):
             if backend == "triton":
                 return TritonTrace(kernel, clients)
-            elif backend == "nki":
-                return NKITrace(kernel, clients)
             else:
                 raise ValueError(f"Unknown backend: {backend}")
 
@@ -174,9 +187,10 @@ def decorator(kernel) -> TraceInterface:
             trace = kernel
             trace.add_client(clients)
             return trace
-        
 
-        raise TypeError(f"Expected JITFunction, InterpretedFunction or Trace, got {type(kernel)}")
+        raise TypeError(
+            f"Expected JITFunction, InterpretedFunction or Trace, got {type(kernel)}"
+        )
 
     return decorator
 
diff --git a/triton_viz/static/gridblock.js b/triton_viz/static/gridblock.js
index 89081726..8e70b823 100644
--- a/triton_viz/static/gridblock.js
+++ b/triton_viz/static/gridblock.js
@@ -2,6 +2,7 @@ import { createMatMulVisualization } from './matmul.js';
 import { createFlipVisualization } from './flip.js';
 import { createLoadVisualization } from './load.js';
 import { createStoreVisualization } from './store.js';
+import { createFlowDiagram } from './nki.js';
 
 export class GridBlock {
     constructor(x, y, width, height, gridX, gridY, gridZ, blockData, onClose, containerElement, canvas, drawFunction) {
@@ -197,12 +198,23 @@ export class GridBlock {
         });
 
         let currentSelectedTab = null;
+        // Tabs for each op
         this.blockData.forEach((op, index) => {
             const opTab = this.createOperationTab(op, index === 0);
             opTab.addEventListener('click', () => this.handleTabClick(opTab, op, currentSelectedTab));
             headerBar.appendChild(opTab);
             if (index === 0) currentSelectedTab = opTab;
         });
+        // Extra: NKI view tab (aggregates all ops)
+        const nkiTab = document.createElement('button');
+        nkiTab.textContent = 'Flow';
+        Object.assign(nkiTab.style, { flex:'0 0 auto', marginRight:'5px', background:'#333', color:'#fff', border:'none', padding:'10px', cursor:'pointer' });
+        nkiTab.addEventListener('click', () => {
+            if (currentSelectedTab) currentSelectedTab.style.backgroundColor = '#333';
+            nkiTab.style.backgroundColor = '#555';
+            this.displayFlowDiagram();
+        });
+        headerBar.appendChild(nkiTab);
 
         return headerBar;
     }
@@ -289,6 +301,14 @@ export class GridBlock {
         }
     }
 
+    displayFlowDiagram() {
+        if (!this.contentArea) return;
+        if (this.visualizationCleanupFunction) { this.visualizationCleanupFunction(); this.visualizationCleanupFunction = null; }
+        this.contentArea.innerHTML = '';
+        // Pass the entire block data (ops in this program) to Flow view
+        this.visualizationCleanupFunction = createFlowDiagram(this.contentArea, this.blockData || []);
+    }
+
 
     createCloseButton() {
         const closeButton = document.createElement('button');
diff --git a/triton_viz/static/load.js b/triton_viz/static/load.js
index 9e08f10d..ba84edb5 100644
--- a/triton_viz/static/load.js
+++ b/triton_viz/static/load.js
@@ -132,6 +132,27 @@ export function createLoadVisualization(containerElement, op) {
         );
 
         addLabels(scene, globalTensor, sliceTensor);
+
+        // Overlay memory flow badges if available (NKI only)
+        try {
+            const badge = document.createElement('div');
+            badge.style.position = 'fixed';
+            badge.style.right = '10px';
+            badge.style.top = '60px';
+            badge.style.zIndex = '2500';
+            badge.style.background = 'rgba(0,0,0,0.65)';
+            badge.style.color = '#fff';
+            badge.style.padding = '6px 8px';
+            badge.style.borderRadius = '6px';
+            badge.style.font = '12px Arial';
+            const ms = (op.mem_src||'').toUpperCase();
+            const md = (op.mem_dst||'').toUpperCase();
+            const by = Number(op.bytes||0);
+            if (ms && md) {
+                badge.innerHTML = `<b>Memory Flow</b><br/>${ms} → ${md}${by?`<br/>${by} B`:''}`;
+                containerElement.appendChild(badge);
+            }
+        } catch(e){}
         const { center } = setupCamera(scene, camera);
         const orbitControls = new OrbitControls(camera, renderer.domElement);
         orbitControls.enableDamping = true;
diff --git a/triton_viz/static/load_utils.js b/triton_viz/static/load_utils.js
index 5a38bc70..80e468c1 100644
--- a/triton_viz/static/load_utils.js
+++ b/triton_viz/static/load_utils.js
@@ -61,9 +61,21 @@ export function createCube(color, tensorName, x, y, z, cubeGeometry, edgesGeomet
 export function createTensor(shape, coords, color, tensorName, cubeGeometry, edgesGeometry, lineMaterial) {
     console.log(`Creating ${tensorName} tensor:`, shape, coords);
     const tensor = new THREE.Group();
-    let [width, height, depth] = shape;
-    depth = depth || 1;
-    height = height || 1;
+    // Normalize shape to width (X), height (Y), depth (Z)
+    let width, height, depth;
+    if (shape.length === 1) {
+        width = shape[0];
+        height = 1;
+        depth = 1;
+    } else if (shape.length === 2) {
+        // Backend provides (H, W) for 2D tensors; interpret as width=W, height=H
+        height = shape[0];
+        width = shape[1];
+        depth = 1;
+    } else {
+        // Assume incoming order already matches [width, height, depth]
+        [width, height, depth] = shape;
+    }
 
     if (tensorName === 'Global') {
         console.log(`Creating global tensor with dimensions: ${width}x${height}x${depth}`);
@@ -158,7 +170,20 @@ export function createTensor(shape, coords, color, tensorName, cubeGeometry, edg
 }
 
 export function calculateTensorSize(shape) {
-    const [width, height, depth] = shape;
+    // Normalize shape for size calculation consistent with createTensor
+    let width, height, depth;
+    if (shape.length === 1) {
+        width = shape[0];
+        height = 1;
+        depth = 1;
+    } else if (shape.length === 2) {
+        // (H, W) -> width=W, height=H
+        height = shape[0];
+        width = shape[1];
+        depth = 1;
+    } else {
+        [width, height, depth] = shape;
+    }
     return new THREE.Vector3(
         width * (CUBE_SIZE + GAP),
         height * (CUBE_SIZE + GAP),
diff --git a/triton_viz/static/matmul.js b/triton_viz/static/matmul.js
index 4db64f8e..d4d89aba 100644
--- a/triton_viz/static/matmul.js
+++ b/triton_viz/static/matmul.js
@@ -19,7 +19,8 @@ export function createMatMulVisualization(containerElement, op) {
 
 
     const sideMenu = document.createElement('div');
-    sideMenu.style.position = 'absolute';
+    // Use fixed position and high z-index to ensure it's above WebGL canvas
+    sideMenu.style.position = 'fixed';
     sideMenu.style.top = '10px';
     sideMenu.style.right = '10px';
     sideMenu.style.width = '200px';
@@ -29,9 +30,31 @@ export function createMatMulVisualization(containerElement, op) {
     sideMenu.style.fontFamily = 'Arial, sans-serif';
     sideMenu.style.fontSize = '14px';
     sideMenu.style.borderRadius = '5px';
-    sideMenu.style.zIndex = '2000';
+    sideMenu.style.zIndex = '3000';
     sideMenu.style.pointerEvents = 'auto';
     containerElement.appendChild(sideMenu);
+
+    // Memory flow badge for Dot (SBUF→PSUM)
+    try {
+        const badge = document.createElement('div');
+        badge.style.position = 'fixed';
+        badge.style.right = '10px';
+        badge.style.top = '60px';
+        badge.style.zIndex = '2500';
+        badge.style.background = 'rgba(0,0,0,0.65)';
+        badge.style.color = '#fff';
+        badge.style.padding = '6px 8px';
+        badge.style.borderRadius = '6px';
+        badge.style.font = '12px Arial';
+        const ms = (op.mem_src||'').toUpperCase();
+        const md = (op.mem_dst||'').toUpperCase();
+        if (ms && md) {
+            const tiles = Array.isArray(op.tile_shape)? `${op.tile_shape[0]}x${op.tile_shape[1]}` : '';
+            const k = Number(op.k||0);
+            badge.innerHTML = `<b>Memory Flow</b><br/>${ms} → ${md}${tiles?`<br/>Tile: ${tiles}`:''}${k?`<br/>K: ${k}`:''}`;
+            containerElement.appendChild(badge);
+        }
+    } catch(e){}
     let hoveredCube = null;
 
 
@@ -50,26 +73,26 @@ export function createMatMulVisualization(containerElement, op) {
     }
 
 
-    function updateSideMenu(matrix, x, y, vectors) {
-        if (!matrix) {
+    function updateSideMenu(matrixOrName, x, y, vectors, value) {
+        if (!matrixOrName) {
             sideMenu.innerHTML = '';
             return;
         }
 
         let matrixName;
         let dims;
-        if (matrix === matrixA) {
-            matrixName = 'A';
-            dims = input_shape;
-        } else if (matrix === matrixB) {
-            matrixName = 'B';
-            dims = other_shape;
-        } else if (matrix === matrixC) {
-            matrixName = 'C';
-            dims = output_shape;
+        // Accept both string name ('A'|'B'|'C') and matrix group object
+        if (typeof matrixOrName === 'string') {
+            const name = matrixOrName.toUpperCase();
+            if (name === 'A') { matrixName = 'A'; dims = input_shape; }
+            else if (name === 'B') { matrixName = 'B'; dims = other_shape; }
+            else if (name === 'C') { matrixName = 'C'; dims = output_shape; }
+            else { sideMenu.innerHTML = ''; return; }
         } else {
-            sideMenu.innerHTML = '';
-            return;
+            if (matrixOrName === matrixA) { matrixName = 'A'; dims = input_shape; }
+            else if (matrixOrName === matrixB) { matrixName = 'B'; dims = other_shape; }
+            else if (matrixOrName === matrixC) { matrixName = 'C'; dims = output_shape; }
+            else { sideMenu.innerHTML = ''; return; }
         }
         console.log(matrixName, "x:", (x + 1), "y:", (y + 1));
         let extra = '';
@@ -90,6 +113,7 @@ export function createMatMulVisualization(containerElement, op) {
             <p>Row: ${y + 1}</p>
             <p>Column: ${x + 1}</p>
             <p>Dimensions: ${dims[0]} x ${dims[1]}</p>
+            <p>Value: ${value !== undefined ? value : 'Loading...'}</p>
             ${extra}
         `;
     }
@@ -152,6 +176,7 @@ export function createMatMulVisualization(containerElement, op) {
 
                 // If hovering C, fetch vectors and update panel
                 let vectors = null;
+                let valueForPanel = res && res.value !== undefined ? res.value : undefined;
                 if (hoveredCube.matrixName === 'C') {
                     try {
                         const resp = await fetch(`${API_BASE}/api/getMatmulVectors`, {
@@ -160,6 +185,15 @@ export function createMatMulVisualization(containerElement, op) {
                             body: JSON.stringify({ uuid: op.uuid, row: hoveredCube.matrixRow, col: hoveredCube.matrixCol })
                         });
                         vectors = await resp.json();
+                        // Compute C[row,col] directly from vectors if available
+                        if (vectors && !vectors.error && Array.isArray(vectors.a_row) && Array.isArray(vectors.b_col)) {
+                            const aRow = vectors.a_row;
+                            const bCol = vectors.b_col;
+                            const len = Math.min(aRow.length, bCol.length);
+                            let sum = 0.0;
+                            for (let i = 0; i < len; i++) sum += aRow[i] * bCol[i];
+                            valueForPanel = sum;
+                        }
                     } catch (e) { vectors = { error: String(e) }; }
 
                     // highlight A[row,:], B[:,col], C[row,col]
@@ -176,7 +210,7 @@ export function createMatMulVisualization(containerElement, op) {
                     // hovering A/B or others -> just reset
                     resetColors();
                 }
-                updateSideMenu(hoveredCube.matrixName, hoveredCube.matrixRow, hoveredCube.matrixCol, vectors);
+                updateSideMenu(hoveredCube.matrixName, hoveredCube.matrixRow, hoveredCube.matrixCol, vectors, valueForPanel);
             }
         } else {
             updateSideMenu(null);
@@ -209,6 +243,10 @@ export function createMatMulVisualization(containerElement, op) {
     renderer.setPixelRatio(dpr);
     renderer.setSize(containerElement.clientWidth, containerElement.clientHeight);
     containerElement.appendChild(renderer.domElement);
+    // Ensure canvas is under overlays
+    renderer.domElement.style.position = 'relative';
+    renderer.domElement.style.zIndex = '1';
+    try { containerElement.style.pointerEvents = 'auto'; } catch(e){}
 
     const ambientLight = new THREE.AmbientLight(0xffffff, 0.5);
     scene.add(ambientLight);
@@ -317,12 +355,12 @@ export function createMatMulVisualization(containerElement, op) {
     }
 
     const controlPanel = document.createElement('div');
-    controlPanel.style.position = 'absolute';
+    controlPanel.style.position = 'fixed';
     controlPanel.style.bottom = '10px';
     controlPanel.style.left = '10px';
     controlPanel.style.display = 'flex';
     controlPanel.style.gap = '10px';
-    controlPanel.style.zIndex = '2000';
+    controlPanel.style.zIndex = '3000';
     controlPanel.style.pointerEvents = 'auto';
 
     // Removed animation controls; keep panel for color-by-value toggle only
diff --git a/triton_viz/static/nki.js b/triton_viz/static/nki.js
new file mode 100644
index 00000000..fd155bc8
--- /dev/null
+++ b/triton_viz/static/nki.js
@@ -0,0 +1,107 @@
+export function createFlowDiagram(containerElement, opsByProgram) {
+// Minimal NKI flow view: three lanes (HBM, SBUF, PSUM) and arrows per op time_idx
+// opsByProgram: array of op objects for a grid block (Load/Store/Dot/Copy) with mem_* fields
+  const laneNames = ["HBM", "SBUF", "PSUM"];
+  const laneY = { HBM: 60, SBUF: 160, PSUM: 260 };
+  const width = containerElement.clientWidth || 1200;
+  const height = 360;
+
+  const canvas = document.createElement('canvas');
+  canvas.width = width;
+  canvas.height = height;
+  canvas.style.width = '100%';
+  canvas.style.height = '360px';
+  canvas.style.background = '#111';
+  canvas.style.border = '1px solid #333';
+  containerElement.appendChild(canvas);
+  const ctx = canvas.getContext('2d');
+
+  // Draw lanes
+  ctx.font = '14px Arial';
+  ctx.fillStyle = '#ddd';
+  ctx.strokeStyle = '#444';
+  laneNames.forEach(name => {
+    const y = laneY[name];
+    ctx.beginPath();
+    ctx.moveTo(80, y);
+    ctx.lineTo(width - 20, y);
+    ctx.stroke();
+    ctx.fillText(name, 20, y + 5);
+  });
+
+  // Normalize time to X range
+  const events = opsByProgram
+    .filter(op => op.time_idx !== undefined && op.time_idx >= 0)
+    .map(op => ({
+      t: op.time_idx,
+      src: (op.mem_src || '').toUpperCase(),
+      dst: (op.mem_dst || '').toUpperCase(),
+      bytes: Number(op.bytes || 0),
+      type: op.type,
+      uuid: op.uuid
+    }));
+  if (events.length === 0) {
+    ctx.fillStyle = '#aaa';
+    ctx.fillText('No NKI metadata found. Run NKI example or add mem_src/mem_dst/bytes/time_idx to records.', 80, height/2);
+    return () => { containerElement.innerHTML = ''; };
+  }
+  events.sort((a,b)=>a.t-b.t);
+  const tMin = events[0].t;
+  const tMax = events[events.length-1].t || (tMin+1);
+  const toX = t => 80 + (width-120) * (t - tMin) / Math.max(1, (tMax - tMin));
+
+  // Color mapping per type
+  const colorFor = (type) => {
+    switch ((type||'').toLowerCase()){
+      case 'load': return '#00bcd4';
+      case 'store': return '#ff9800';
+      case 'dot': return '#8bc34a';
+      default: return '#9e9e9e';
+    }
+  };
+
+  // Draw arrows for each event
+  events.forEach(e => {
+    const y1 = laneY[e.src] ?? laneY.SBUF;
+    const y2 = laneY[e.dst] ?? laneY.SBUF;
+    const x = toX(e.t);
+    const col = colorFor(e.type);
+    // Arrow thickness by bytes (log scaled)
+    const w = Math.max(1, Math.log10((e.bytes||1)));
+    ctx.strokeStyle = col;
+    ctx.lineWidth = w;
+    ctx.beginPath();
+    ctx.moveTo(x, y1);
+    ctx.lineTo(x, y2);
+    ctx.stroke();
+    // Arrow head
+    ctx.beginPath();
+    const dir = Math.sign(y2 - y1) || 1;
+    const hx = x; const hy = y2;
+    ctx.moveTo(hx, hy);
+    ctx.lineTo(hx - 6, hy - 6*dir);
+    ctx.lineTo(hx + 6, hy - 6*dir);
+    ctx.closePath();
+    ctx.fillStyle = col;
+    ctx.fill();
+    // Label
+    ctx.fillStyle = '#ccc';
+    ctx.font = '11px Arial';
+    const label = `${e.type}  ${e.bytes||0}B`;
+    ctx.fillText(label, x + 6, (y1+y2)/2 - 6);
+  });
+
+  // Legend
+  const legend = document.createElement('div');
+  legend.style.position = 'relative';
+  legend.style.marginTop = '6px';
+  legend.style.color = '#ddd';
+  legend.style.font = '12px Arial';
+  legend.innerHTML = `<b>Flow Diagram</b> &nbsp; <span style="color:#00bcd4">■ Load</span>&nbsp;&nbsp;<span style="color:#8bc34a">■ Dot(PSUM)</span>&nbsp;&nbsp;<span style="color:#ff9800">■ Store</span>`;
+  containerElement.appendChild(legend);
+
+  return () => { containerElement.innerHTML = ''; };
+}
+
+// Backward compatibility alias
+export const createNKIFlow = createFlowDiagram;
diff --git a/triton_viz/static/store.js b/triton_viz/static/store.js
index cfca8838..b276c083 100644
--- a/triton_viz/static/store.js
+++ b/triton_viz/static/store.js
@@ -70,6 +70,27 @@ export function createStoreVisualization(containerElement, op) {
         scene.add(sliceTensor);
 
         addLabels(scene, globalTensor, sliceTensor);
+
+        // Overlay memory flow badges if available (NKI only)
+        try {
+            const badge = document.createElement('div');
+            badge.style.position = 'fixed';
+            badge.style.right = '10px';
+            badge.style.top = '60px';
+            badge.style.zIndex = '2500';
+            badge.style.background = 'rgba(0,0,0,0.65)';
+            badge.style.color = '#fff';
+            badge.style.padding = '6px 8px';
+            badge.style.borderRadius = '6px';
+            badge.style.font = '12px Arial';
+            const ms = (op.mem_src||'').toUpperCase();
+            const md = (op.mem_dst||'').toUpperCase();
+            const by = Number(op.bytes||0);
+            if (ms && md) {
+                badge.innerHTML = `<b>Memory Flow</b><br/>${ms} → ${md}${by?`<br/>${by} B`:''}`;
+                containerElement.appendChild(badge);
+            }
+        } catch(e){}
         const { center } = setupCamera(scene, camera);
         const orbitControls = new OrbitControls(camera, renderer.domElement);
         orbitControls.enableDamping = true;
diff --git a/triton_viz/visualizer/draw.py b/triton_viz/visualizer/draw.py
index 6ea66c58..90289230 100644
--- a/triton_viz/visualizer/draw.py
+++ b/triton_viz/visualizer/draw.py
@@ -90,9 +90,10 @@ def extract_load_coords(
         record.masks,
     )
 
+    # Report (x, y, z); delinearized returns (z, y, x)
     global_coords = [
         (float(xi), float(yi), float(zi))
-        for xi, yi, zi in zip(global_z, global_y, global_x)
+        for xi, yi, zi in zip(global_x, global_y, global_z)
         if xi != -1 and yi != -1 and zi != -1
     ]
 
@@ -172,6 +173,12 @@ def prepare_visualization_data(program_records, tensor_table):
                     "uuid": record_uuid,
                     # provide C values for color-by-value
                     "c_shape": record.output_shape,
+                    # NKI meta (accumulate to PSUM)
+                    "mem_src": getattr(record, "mem_src", None),
+                    "mem_dst": getattr(record, "mem_dst", None),
+                    "tile_shape": getattr(record, "tile_shape", None),
+                    "k": int(getattr(record, "k", 0)),
+                    "time_idx": int(getattr(record, "time_idx", -1)),
                 }
             )
 
@@ -238,12 +245,46 @@ def prepare_visualization_data(program_records, tensor_table):
                     "global_coords": global_coords,
                     "slice_coords": slice_coords,
                     "uuid": record_uuid,
+                    # NKI flow meta (optional)
+                    "mem_src": getattr(record, "mem_src", None),
+                    "mem_dst": getattr(record, "mem_dst", None),
+                    "bytes": int(getattr(record, "bytes", 0)),
+                    "time_idx": int(getattr(record, "time_idx", -1)),
                 }
             )
 
+            # Normalize to torch.Tensor for downstream APIs
+            try:
+                import numpy as _np
+                import torch as _torch
+
+                gt = global_tensor.data
+                # gt could be torch.Tensor, numpy.ndarray, or NDArray wrapper
+                if (
+                    hasattr(gt, "cpu")
+                    and callable(getattr(gt, "cpu"))
+                    and hasattr(gt, "shape")
+                ):
+                    t_cpu = gt.cpu()
+                else:
+                    # NDArray: prefer .data property -> numpy array
+                    if hasattr(gt, "data"):
+                        arr = gt.data
+                    else:
+                        arr = gt
+                    arr = _np.asarray(arr)
+                    t_cpu = _torch.from_numpy(
+                        arr.copy() if not arr.flags.c_contiguous else arr
+                    )
+            except Exception:
+                # Fallback to empty tensor if conversion fails
+                import torch as _torch
+
+                t_cpu = _torch.tensor([])
+
             raw_tensor_data[record_uuid] = {
-                "global_tensor": global_tensor.data.cpu(),  # Ensure it's on CPU
-                "dims": len(global_tensor.data.cpu().shape),
+                "global_tensor": t_cpu,
+                "dims": len(t_cpu.shape),
                 "tracebacks": [
                     {
                         "filename": f.filename,
@@ -269,6 +310,10 @@ def prepare_visualization_data(program_records, tensor_table):
                     "global_coords": global_coords,
                     "slice_coords": slice_coords,
                     "uuid": record_uuid,
+                    "mem_src": getattr(record, "mem_src", None),
+                    "mem_dst": getattr(record, "mem_dst", None),
+                    "bytes": int(getattr(record, "bytes", 0)),
+                    "time_idx": int(getattr(record, "time_idx", -1)),
                 }
             )
 
diff --git a/triton_viz/visualizer/interface.py b/triton_viz/visualizer/interface.py
index c287b5ef..52156d2b 100644
--- a/triton_viz/visualizer/interface.py
+++ b/triton_viz/visualizer/interface.py
@@ -394,13 +394,30 @@ def get_load_value():
         x is not None and y is not None and z is not None
     ):
         try:
+            import numpy as _np
+
+            t = op_data["global_tensor"]
+            # Normalize to numpy array for robust indexing
+            if hasattr(t, "cpu"):
+                try:
+                    arr = t.cpu().numpy()
+                except Exception:
+                    arr = _np.asarray(t)
+            elif hasattr(t, "_value"):
+                arr = _np.asarray(getattr(t, "_value"))
+            elif hasattr(t, "data"):
+                arr = _np.asarray(getattr(t, "data"))
+            else:
+                arr = _np.asarray(t)
+
+            yy, xx, zz = int(y), int(x), int(z)
             value = 0.0
-            if op_data["dims"] == 3:
-                value = op_data["global_tensor"][x, y, z].item()
-            elif op_data["dims"] == 2:
-                value = op_data["global_tensor"][x, y].item()
-            elif op_data["dims"] == 1:
-                value = op_data["global_tensor"][x].item()
+            if arr.ndim >= 3:
+                value = float(arr[yy, xx, zz])
+            elif arr.ndim == 2:
+                value = float(arr[yy, xx])
+            elif arr.ndim == 1:
+                value = float(arr[xx])
 
             return jsonify({"value": value})
         except IndexError:
@@ -470,24 +487,34 @@ def get_load_tensor():
     if "global_tensor" not in op_data:
         return jsonify({"error": "Global tensor data not found"}), 200
 
-    t = op_data["global_tensor"].cpu()
-    try:
-        t_min = float(t.min().item())
-        t_max = float(t.max().item())
-    except Exception:
-        # In case of empty tensor
-        t_min = 0.0
-        t_max = 0.0
+    import numpy as _np
 
-    return jsonify(
-        {
-            "shape": list(t.shape),
-            "dims": len(t.shape),
-            "min": t_min,
-            "max": t_max,
-            "values": t.numpy().tolist(),
-        }
-    )
+    try:
+        t = op_data["global_tensor"]
+        if hasattr(t, "cpu"):
+            try:
+                arr = t.cpu().numpy()
+            except Exception:
+                arr = _np.asarray(t)
+        elif hasattr(t, "_value"):
+            arr = _np.asarray(getattr(t, "_value"))
+        elif hasattr(t, "data"):
+            arr = _np.asarray(getattr(t, "data"))
+        else:
+            arr = _np.asarray(t)
+        t_min = float(_np.min(arr)) if arr.size else 0.0
+        t_max = float(_np.max(arr)) if arr.size else 0.0
+        return jsonify(
+            {
+                "shape": list(arr.shape),
+                "dims": int(arr.ndim),
+                "min": t_min,
+                "max": t_max,
+                "values": arr.tolist(),
+            }
+        )
+    except Exception as e:
+        return jsonify({"error": f"getLoadTensor failed: {e}"}), 200
 
 
 def run_flask_with_cloudflared(port: int = 8000, tunnel_port: int | None = None):

From b959a125c9a0315af38d9ce71a6a5a3a30ea7f11 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Tue, 28 Oct 2025 16:30:25 -0400
Subject: [PATCH 57/90] patch nki pt.1

---
 triton_viz/core/nki.py   | 25 ++++++++++++++++++-------
 triton_viz/core/trace.py |  9 +++++----
 2 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index 048b802d..1314fe17 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -107,6 +107,7 @@ def __init__(self, buffer=None, name="", **kwargs):
         if "value" in kwargs:
             assert val is None or val.shape == kwargs["value"].shape
             val = kwargs["value"]
+        self._data_ptr = None
         self.data = val
 
     @property
@@ -118,7 +119,9 @@ def dtype(self):
         return self.data.dtype if self.data is not None else None
 
     def data_ptr(self):
-        return self.data.ctypes.data
+        if self._data_ptr is None:
+            self._data_ptr = self.data.ctypes.data
+        return self._data_ptr
 
     def stride(self):
         return self.data.strides
@@ -562,9 +565,6 @@ def run(self, *args, **kwargs):
 
         patch()  # NKI interpreter patching
 
-        # with client_manager.patch():
-        # v
-
         # Apply AST transformer to convert nl.load/nl.store calls to nl.masked_load/nl.masked_store
         if hasattr(self.fn, "__code__"):
             # Get the source code of the function
@@ -590,6 +590,14 @@ def run(self, *args, **kwargs):
 
         # convert args to NDArray if they are not already
         args = [arg if isinstance(arg, NDArray) else NDArray(value=arg) for arg in args]
+
+        name_args = inspect.getcallargs(self.fn, *args)
+        call_args = {}
+        for name, arg in name_args.items():
+            call_args[name] = arg
+            ret = arg
+            client_manager.arg_callback(name, arg, ret)
+
         for x in range(grid_dims[0]):
             for y in range(grid_dims[1]):
                 for z in range(grid_dims[2]):
@@ -599,8 +607,11 @@ def run(self, *args, **kwargs):
                     if client_manager is not None:
                         client_manager.grid_idx_callback((x, y, z))
 
-                    result = self.fn(*args, **kwargs)
-        # ^
+                    if not client_manager.pre_run_callback(self.fn):
+                        return
+                    self.fn(*args, **kwargs)
+                    if not client_manager.post_run_callback(self.fn):
+                        return
 
         unpatch()
-        return result.data
+        # return result.data
diff --git a/triton_viz/core/trace.py b/triton_viz/core/trace.py
index ac6a4467..1789bc57 100644
--- a/triton_viz/core/trace.py
+++ b/triton_viz/core/trace.py
@@ -145,10 +145,11 @@ def __call__(self, *args, **kwargs):
         return self[(1, 1, 1)](*args, **kwargs)
 
     def run(self, *args, **kwargs):
-        kwargs.update({"client_manager": self.client_manager})
-        ret = self.interpreter_fn.run(*args, **kwargs)
-        self.finalize()
-        return ret
+        with self.client_manager.patch_run(self.interpreter_fn):
+            kwargs.update({"client_manager": self.client_manager})
+            ret = self.interpreter_fn.run(*args, **kwargs)
+            self.finalize()
+            return ret
 
 
 def trace(clients: Union[str, Client, None] = None, backend: str = "triton"):

From ed0b309748a0fc292cb683664f71f758e594ccac Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Tue, 28 Oct 2025 16:36:16 -0400
Subject: [PATCH 58/90] nki patch pt.2 (new tracer callbacks)

---
 triton_viz/clients/tracer/tracer.py | 65 +++++++++++++++++++++++++----
 triton_viz/core/data.py             |  6 +++
 triton_viz/core/patch.py            | 50 ++++++++++++----------
 3 files changed, 91 insertions(+), 30 deletions(-)

diff --git a/triton_viz/clients/tracer/tracer.py b/triton_viz/clients/tracer/tracer.py
index 57005b33..c50fa0ea 100644
--- a/triton_viz/clients/tracer/tracer.py
+++ b/triton_viz/clients/tracer/tracer.py
@@ -1,6 +1,7 @@
 from ...core.client import Client
 from ...core.callbacks import OpCallbacks, ForLoopCallbacks
-from ...core.data import Op, Load, Store, ReduceSum, Dot, Grid, RawLoad, RawStore
+from ...core.data import Op, Load, Store, ReduceSum, Dot, Grid, RawLoad, RawStore, Array
+from triton_viz.core.nki_masked_load import masked_load
 from typing import Callable, Optional, Union
 import numpy as np
 
@@ -70,10 +71,8 @@ def grid_idx_callback(self, grid_idx: tuple[int, ...]):
     def grid_callback(self, grid: tuple[int, ...]):
         self.tensors = sorted(self.tensors, key=lambda x: x.data_ptr())
 
-    def register_op_callback(self, op_type: type[Op], *ignore_args, **ignore_kwargs) -> OpCallbacks:
-        def pre_load_callback(
-            ptr, mask, *ignore_args, **ignore_kwargs
-        ):
+    def register_op_callback(self, op_type: type[Op]) -> OpCallbacks:
+        def pre_load_callback(ptr, mask, *ignore_args, **ignore_kwargs):
             if not self.sample:
                 return
             first_ptr = np.reshape(ptr.data, (-1))[0]
@@ -82,6 +81,34 @@ def pre_load_callback(
                 Load(tensor.data_ptr(), ptr.data - tensor.data_ptr(), mask.data)
             )
 
+        def _convert_keys_to_numpy(keys):
+            """Convert any NDArrays in keys to numpy arrays."""
+            if isinstance(keys, (tuple, list)):
+                return tuple(_convert_keys_to_numpy(k) for k in keys)
+            elif hasattr(keys, "data"):
+                return keys.data
+            else:
+                return keys
+
+        def post_array_callback(ret, *ignore_args, **ignore_kwargs):
+            assert hasattr(ret, "data")
+            self.tensors.append(ret)
+
+        def pre_masked_load_callback(
+            ptr, keys, mask=None, *ignore_args, **ignore_kwargs
+        ):
+            if not self.sample:
+                return
+            keys = _convert_keys_to_numpy(keys)
+
+            self.records.append(
+                Load(
+                    ptr.data_ptr(),
+                    masked_load(ptr.get_offsets().data, keys, mask=mask.data),
+                    mask.data,
+                )
+            )
+
         def pre_store_callback(ptr, value, mask, *ignore_args, **ignore_kwargs):
             if not self.sample:
                 return
@@ -91,6 +118,20 @@ def pre_store_callback(ptr, value, mask, *ignore_args, **ignore_kwargs):
                 Store(tensor.data_ptr(), ptr.data - tensor.data_ptr(), mask.data)
             )
 
+        def pre_masked_store_callback(
+            ptr, keys, value, mask=None, *ignore_args, **ignore_kwargs
+        ):
+            if not self.sample:
+                return
+            keys = _convert_keys_to_numpy(keys)
+            offsets = masked_load(ptr.get_offsets().data, keys)
+            if mask is None:
+                mask_data = np.ones_like(offsets).astype(bool)
+            else:
+                mask_data = mask.data
+
+            self.records.append(Store(ptr.data_ptr(), offsets, mask_data))
+
         # Raw (unmasked) ops: synthesize a full True mask based on ptr shape
         def pre_raw_load_callback(ptr):
             if not self.sample:
@@ -110,7 +151,9 @@ def pre_raw_store_callback(ptr, value):
             true_mask = np.ones_like(offsets, dtype=bool)
             self.records.append(Store(tensor.data_ptr(), offsets, true_mask))
 
-        def post_reduce_sum_callback(ret, input, axis=None, keep_dims=False, *ignore_args, **ignore_kwargs):
+        def post_reduce_sum_callback(
+            ret, input, axis=None, keep_dims=False, *ignore_args, **ignore_kwargs
+        ):
             if not self.sample:
                 return
             input_shape = input.handle.data.shape
@@ -123,12 +166,20 @@ def post_dot_callback(ret, input, other, *ignore_args, **ignore_kwargs):
             input_shape = input.data.shape
             other_shape = other.data.shape
             ret_shape = ret.data.shape
-            self.records.append(Dot(input_shape, other_shape, ret_shape))
+            self.records.append(
+                Dot(input_shape, other_shape, ret_shape, input.data, other.data)
+            )
 
+        if op_type is Array:  # THTODO: only for NKI
+            return OpCallbacks(after_callback=post_array_callback)
         if op_type is Load:
             return OpCallbacks(before_callback=pre_load_callback)
         elif op_type is Store:
             return OpCallbacks(before_callback=pre_store_callback)
+        # if op_type is Load: # THTODO: only for NKI
+        #    return OpCallbacks(before_callback=pre_masked_load_callback)
+        # elif op_type is Store: # THTODO: only for NKI
+        #    return OpCallbacks(before_callback=pre_masked_store_callback)
         elif op_type is RawLoad:
             return OpCallbacks(before_callback=pre_raw_load_callback)
         elif op_type is RawStore:
diff --git a/triton_viz/core/data.py b/triton_viz/core/data.py
index 41182048..cca41123 100644
--- a/triton_viz/core/data.py
+++ b/triton_viz/core/data.py
@@ -34,6 +34,12 @@ class ProgramId(Op):
     name: ClassVar[str] = "program_id"
 
 
+@dataclass
+class Array(Op):
+    name: ClassVar[str] = "array"
+    ptr: int
+
+
 @dataclass
 class RawStore(Op):
     name: ClassVar[str] = "raw_store"
diff --git a/triton_viz/core/patch.py b/triton_viz/core/patch.py
index de1f7cac..8f83d623 100644
--- a/triton_viz/core/patch.py
+++ b/triton_viz/core/patch.py
@@ -8,6 +8,7 @@
 from .callbacks import OpCallbacks, ForLoopCallbacks
 from .data import (
     Op,
+    Array,
     RawLoad,
     Load,
     RawStore,
@@ -159,33 +160,36 @@
         AtomicCas: "create_atomic_cas",
     }
 elif BUILDER == nki_builder:
-    op_list = [
-        ProgramId,
-        Store,
-        Load,
-        Dot,
-        UnaryOp,
-        MakeRange
-    ]
+    op_list = [Array, ProgramId, Store, Load, Dot, UnaryOp, MakeRange]
     original_ops = {
         ProgramId: nki_builder.program_id,
-        Store: nki_builder.store,
-        Load: nki_builder.load,
+        Array: nki_builder.ndarray,
+        Load: nki_builder.masked_load,
+        Store: nki_builder.masked_store,
         Dot: nki_builder.matmul,
-        UnaryOp: nki_builder.unary_op,
-        #BinaryOp: nki_builder.binary_op,
-        #TernaryOp: nki_builder.ternary_op,
+        UnaryOp: nki_builder._unary_op,
+        # BinaryOp: nki_builder.binary_op,
+        # TernaryOp: nki_builder.ternary_op,
         MakeRange: nki_builder.arange,
-        #AddPtr: nki_builder.create_addptr,
-        #ExpandDims: nki_builder.create_expand_dims,
-        #Broadcast: nki_builder.create_broadcast,
-        #Splat: nki_builder.create_splat,
-        #MakeBlockPointer: nki_builder.create_make_block_ptr,
-        #TensorPointerLoad: nki_builder.create_tensor_pointer_load,
-        #TensorPointerStore: nki_builder.create_tensor_pointer_store,
-        #Idiv: nki_builder.create_idiv,
-        #Rsqrt: nki_builder.create_rsqrt,
-        #CastImpl: nki_builder.cast_impl,
+        # AddPtr: nki_builder.create_addptr,
+        # ExpandDims: nki_builder.create_expand_dims,
+        # Broadcast: nki_builder.create_broadcast,
+        # Splat: nki_builder.create_splat,
+        # MakeBlockPointer: nki_builder.create_make_block_ptr,
+        # TensorPointerLoad: nki_builder.create_tensor_pointer_load,
+        # TensorPointerStore: nki_builder.create_tensor_pointer_store,
+        # Idiv: nki_builder.create_idiv,
+        # Rsqrt: nki_builder.create_rsqrt,
+        # CastImpl: nki_builder.cast_impl,
+    }
+    _OP_ATTR_NAMES = {
+        ProgramId: "program_id",
+        Array: "ndarray",
+        Load: "masked_load",
+        Store: "masked_store",
+        Dot: "matmul",
+        UnaryOp: "_unary_op",
+        MakeRange: "arange",
     }
 
 reduce_map: dict[type[Op], Callable] = {

From 2e1943e57160747e346253860246c5cc705573cc Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Tue, 28 Oct 2025 16:37:27 -0400
Subject: [PATCH 59/90] nki patch pt.3 (nki-specific changes to allow nki
 tracing)

---
 triton_viz/clients/tracer/tracer.py | 16 ++++++++--------
 triton_viz/core/client.py           |  6 +++---
 triton_viz/core/patch.py            | 18 +++++++++---------
 3 files changed, 20 insertions(+), 20 deletions(-)

diff --git a/triton_viz/clients/tracer/tracer.py b/triton_viz/clients/tracer/tracer.py
index c50fa0ea..d7d268d1 100644
--- a/triton_viz/clients/tracer/tracer.py
+++ b/triton_viz/clients/tracer/tracer.py
@@ -172,14 +172,14 @@ def post_dot_callback(ret, input, other, *ignore_args, **ignore_kwargs):
 
         if op_type is Array:  # THTODO: only for NKI
             return OpCallbacks(after_callback=post_array_callback)
-        if op_type is Load:
-            return OpCallbacks(before_callback=pre_load_callback)
-        elif op_type is Store:
-            return OpCallbacks(before_callback=pre_store_callback)
-        # if op_type is Load: # THTODO: only for NKI
-        #    return OpCallbacks(before_callback=pre_masked_load_callback)
-        # elif op_type is Store: # THTODO: only for NKI
-        #    return OpCallbacks(before_callback=pre_masked_store_callback)
+        # if op_type is Load:
+        #     return OpCallbacks(before_callback=pre_load_callback)
+        # elif op_type is Store:
+        #     return OpCallbacks(before_callback=pre_store_callback)
+        if op_type is Load:  # THTODO: only for NKI
+            return OpCallbacks(before_callback=pre_masked_load_callback)
+        elif op_type is Store:  # THTODO: only for NKI
+            return OpCallbacks(before_callback=pre_masked_store_callback)
         elif op_type is RawLoad:
             return OpCallbacks(before_callback=pre_raw_load_callback)
         elif op_type is RawStore:
diff --git a/triton_viz/core/client.py b/triton_viz/core/client.py
index c5d2c52c..8eadee4d 100644
--- a/triton_viz/core/client.py
+++ b/triton_viz/core/client.py
@@ -15,7 +15,7 @@
 )
 from functools import wraps
 from .callbacks import OpCallbacks, ForLoopCallbacks
-from .patch import patch_lang, unpatch_lang
+# from .patch import patch_lang, unpatch_lang
 
 
 class Client(ABC):
@@ -138,14 +138,14 @@ def patch_run(self, fn):
                 loop_callbacks = client.register_for_loop_callback()
                 patch_for_loop(loop_callbacks)
                 # Remaps core language functions to interpreted ones
-                patch_lang(fn)
+                # patch_lang(fn)
             try:
                 yield
             finally:
                 for op in op_list:
                     unpatch_op(op)
                 unpatch_for_loop()
-                unpatch_lang()
+                # unpatch_lang()
 
     def pre_run_callback(self, fn: Callable) -> bool:
         rets = []
diff --git a/triton_viz/core/patch.py b/triton_viz/core/patch.py
index 8f83d623..3d342598 100644
--- a/triton_viz/core/patch.py
+++ b/triton_viz/core/patch.py
@@ -56,7 +56,7 @@
 from triton.runtime import JITFunction
 from triton_viz.core.nki import nki_builder
 
-BUILDER = interpreter_builder
+BUILDER = nki_builder
 if BUILDER == interpreter_builder:
     op_list = [
         ProgramId,
@@ -171,17 +171,17 @@
         # BinaryOp: nki_builder.binary_op,
         # TernaryOp: nki_builder.ternary_op,
         MakeRange: nki_builder.arange,
-        # AddPtr: nki_builder.create_addptr,
-        # ExpandDims: nki_builder.create_expand_dims,
-        # Broadcast: nki_builder.create_broadcast,
-        # Splat: nki_builder.create_splat,
-        # MakeBlockPointer: nki_builder.create_make_block_ptr,
-        # TensorPointerLoad: nki_builder.create_tensor_pointer_load,
-        # TensorPointerStore: nki_builder.create_tensor_pointer_store,
-        # Idiv: nki_builder.create_idiv,
         # Rsqrt: nki_builder.create_rsqrt,
         # CastImpl: nki_builder.cast_impl,
     }
+    _OP_ATTR_NAMES = {
+        ProgramId: "program_id",
+        Array: "ndarray",
+        Load: "masked_load",
+        Store: "masked_store",
+        Dot: "matmul",
+        UnaryOp: "_unary_op",
+    }
     _OP_ATTR_NAMES = {
         ProgramId: "program_id",
         Array: "ndarray",

From 301aa66e711d57b3027b2c690e96f47785c6f8e9 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Tue, 28 Oct 2025 16:38:13 -0400
Subject: [PATCH 60/90] nki patch pt.4 (update matmul example)

---
 nki-examples/matmul.py | 47 +++++++-----------------------------------
 1 file changed, 8 insertions(+), 39 deletions(-)

diff --git a/nki-examples/matmul.py b/nki-examples/matmul.py
index 2a8b8edc..2b72de57 100644
--- a/nki-examples/matmul.py
+++ b/nki-examples/matmul.py
@@ -6,11 +6,8 @@
 from triton_viz.core.trace import launches
 import numpy as np
 
-tensors = set()  # THTODO: rm
-records = [triton_viz.core.data.Grid((0, 0, 0))]  # THTODO: rm
 
-
-def matmul_kernel(lhs, rhs):
+def matmul_kernel(lhs, rhs, result):
     """NKI matmul_kernel to compute a matrix multiplication operation in a tiled manner
 
     Args:
@@ -27,15 +24,11 @@ def matmul_kernel(lhs, rhs):
     M, K = lhs.shape
     K_, N = rhs.shape
     assert K == K_, "lhs and rhs must have the same contraction dimension"
-    result = nl.ndarray((M, N), dtype=lhs.dtype, buffer=nl.shared_hbm)
 
     TILE_M = 2
     TILE_K = 2
     TILE_N = 4
 
-    tensors.add(lhs)
-    tensors.add(rhs)
-    tensors.add(result)
     # Use affine_range to loop over tiles
     for m in nl.affine_range(M // TILE_M):
         for n in nl.affine_range(N // TILE_N):
@@ -52,35 +45,15 @@ def matmul_kernel(lhs, rhs):
                 lhs_f = nl.arange(TILE_K)[None, :] + k * TILE_K
                 lhs_mask = (lhs_p < M) & (lhs_f < K)
                 lhs_tile = nl.load(lhs[lhs_p, lhs_f], mask=lhs_mask)
-                records.append(
-                    triton_viz.core.data.Load(
-                        lhs.data_ptr(), (4 * (lhs_p * K + lhs_f)).data, lhs_mask.data
-                    )
-                )  # THTODO: rm
 
                 rhs_p = nl.arange(TILE_K)[:, None] + k * TILE_K
                 rhs_f = nl.arange(TILE_N)[None, :] + n * TILE_N
                 rhs_mask = (rhs_p < K) & (rhs_f < N)
                 rhs_tile = nl.load(rhs[rhs_p, rhs_f], mask=rhs_mask)
-                records.append(
-                    triton_viz.core.data.Load(
-                        rhs.data_ptr(), (4 * (rhs_p * N + rhs_f)).data, rhs_mask.data
-                    )
-                )  # THTODO: rm
 
                 # Accumulate partial-sums into PSUM
                 x = nl.matmul(lhs_tile[...], rhs_tile[...], transpose_x=False)
                 res_psum += x
-                records.append(
-                    triton_viz.core.data.Dot(
-                        lhs_tile.shape,
-                        rhs_tile.shape,
-                        res_psum.shape,
-                        lhs_tile.data,
-                        rhs_tile.data,
-                        x.data,
-                    )
-                )  # THTODO: rm
 
             # Copy the result from PSUM back to SBUF, and cast to expected output data-type
             res_sb = nl.copy(res_psum, dtype=result.dtype)
@@ -91,12 +64,8 @@ def matmul_kernel(lhs, rhs):
             nl.store(
                 result[m * TILE_M : (m + 1) * TILE_M, n * TILE_N : (n + 1) * TILE_N],
                 value=res_sb,
+                mask=out_mask,
             )
-            records.append(
-                triton_viz.core.data.Store(
-                    result.data_ptr(), (4 * (out_p * N + out_f)).data, out_mask.data
-                )
-            )  # THTODO: rm
 
     return result
 
@@ -105,19 +74,18 @@ def matmul_kernel(lhs, rhs):
 kernel_grid = (1, 1, 1)
 lhs_small = np.arange(16).astype(np.float32).reshape(4, 4)
 rhs_small = np.arange(32).astype(np.float32).reshape(4, 8)
-kernel_args = (lhs_small, rhs_small)
+result = np.empty((lhs_small.shape[0], rhs_small.shape[1]), dtype=lhs_small.dtype)
+kernel_args = (lhs_small, rhs_small, result)
 
 if TRITON_VIZ:
     print("Executing matmul_kernel with NKI interpreter...")
     traced_kernel = triton_viz.trace(clients=Tracer(), backend="nki")(matmul_kernel)
-    kk = traced_kernel[kernel_grid]
-    z2 = kk(*kernel_args)
+    kernel = traced_kernel[kernel_grid]
+    kernel(*kernel_args)
 
     print(f"Number of launches: {len(launches)}")
     if launches:
         launch = launches[-1]
-        launch.records = records  # THTODO: rm
-        launch.tensors = tensors  # THTODO: rm
         print(f"Number of records: {len(launch.records)}")
         for i, record in enumerate(launch.records):
             print(f"Record {i}: {type(record).__name__}")
@@ -139,8 +107,9 @@ def matmul_kernel(lhs, rhs):
 else:
     print("Executing NKI JIT-ed matmul_kernel...")
     compiled_kernel = nki.jit(matmul_kernel)
-    z2 = nki.simulate_kernel(compiled_kernel[kernel_grid], *kernel_args)
+    nki.simulate_kernel(compiled_kernel[kernel_grid], *kernel_args)
 
+z2 = result
 z1 = lhs_small @ rhs_small
 print(np.max(np.abs(z1 - z2)))
 assert np.allclose(z1, z2)

From e9ecf825a64a0a20ddf58d852f2605547d3f2807 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Tue, 28 Oct 2025 18:13:19 -0400
Subject: [PATCH 61/90] remove NLSlice (was never used)

---
 triton_viz/core/nki.py | 104 +----------------------------------------
 1 file changed, 1 insertion(+), 103 deletions(-)

diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index 3cb3e84a..abb2ab1c 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -6,94 +6,6 @@
 from .nki_masked_load import masked_load, masked_store
 
 
-# Q1: slicing semantic is weird
-# Q2: why cannot we execute the .func function?
-
-
-# Multi-dimensional slice class
-class NLSlice:
-    def __init__(self, start=None, stop=None, step: int = 1):
-        self.start = [start] if isinstance(start, int) else start
-        self.stop = [stop] if isinstance(stop, int) else stop
-        self.step = [step] if isinstance(step, int) else step
-
-    def __repr__(self):
-        repr = ""
-        for start, stop, step in zip(self.start, self.stop, self.step):
-            if start is None:
-                start = "None"
-            if stop is None:
-                stop = "None"
-            if step is None:
-                step = "None"
-            repr += f"(start={start}, stop={stop}, step={step}) "
-        return repr
-
-    def __add__(self, other):
-        new_start = []
-        new_stop = []
-        new_step = []
-        if isinstance(other, NLSlice):
-            for start, stop, step in zip(self.start, self.stop, self.step):
-                new_start.append(start + other.start if start is not None else None)
-                new_stop.append(stop + other.stop if stop is not None else None)
-                new_step.append(step + other.step if step is not None else None)
-            return NLSlice(start=new_start, stop=new_stop, step=new_step)
-        elif isinstance(other, int):
-            for start, stop, step in zip(self.start, self.stop, self.step):
-                new_start.append(start + other if start is not None else None)
-                new_stop.append(stop + other if stop is not None else None)
-                new_step.append(step)
-            return NLSlice(start=new_start, stop=new_stop, step=new_step)
-        raise TypeError(
-            f"Unsupported operand type(s) for +: 'NLSlice' and '{type(other).__name__}'"
-        )
-
-    def __radd__(self, other):
-        new_start = []
-        new_stop = []
-        new_step = []
-        if isinstance(other, NLSlice):
-            for start, stop, step in zip(self.start, self.stop, self.step):
-                new_start.append(other.start + start if start is not None else None)
-                new_stop.append(other.stop + stop if stop is not None else None)
-                new_step.append(other.step + step if step is not None else None)
-            return NLSlice(start=new_start, stop=new_stop, step=new_step)
-        elif isinstance(other, int):
-            for start, stop, step in zip(self.start, self.stop, self.step):
-                new_start.append(other + start if start is not None else None)
-                new_stop.append(other + stop if stop is not None else None)
-                new_step.append(step)
-            return NLSlice(start=new_start, stop=new_stop, step=new_step)
-        raise TypeError(
-            f"Unsupported operand type(s) for +: '{type(other).__name__}' and 'NLSlice'"
-        )
-
-    def __getitem__(self, keys):
-        new_start = []
-        new_stop = []
-        new_step = []
-        idx = 0
-        for k in keys:
-            # check if k is None:
-            if k is None:
-                new_start.append(None)
-                new_stop.append(None)
-                new_step.append(None)
-            elif isinstance(k, slice):
-                assert (
-                    k.start is None and k.stop is None and k.step is None
-                ), "Slice must be complete"
-                new_start.append(self.start[idx])
-                new_stop.append(self.stop[idx])
-                new_step.append(self.step[idx])
-                idx += 1
-            else:
-                raise TypeError(f"Unsupported key type: {type(k)}")
-
-        return NLSlice(start=new_start, stop=new_stop, step=new_step)
-
-
 class NDArray:
     def __init__(self, buffer=None, name="", **kwargs):
         self.buffer = buffer
@@ -183,21 +95,7 @@ def __getitem__(self, keys):
             keys = (keys,)
 
         # Apply the slicing to the underlying numpy array
-        new_keys = []
-        arr_dim = 0
-        for k in keys:
-            if isinstance(k, NDArray):
-                dim_len = self.data.shape[arr_dim]
-                new_keys.append(k.data.clip(0, dim_len - 1))
-            elif isinstance(k, NLSlice):
-                new_keys.append(slice(k.start, k.stop, k.step))
-            elif k is None:
-                new_keys.append(k)
-                arr_dim -= 1  # add new dim -> revisit arr_dim for next key
-            else:
-                new_keys.append(k)
-            arr_dim += 1
-
+        new_keys = [k.data if isinstance(k, NDArray) else k for k in keys]
         sliced_value = self.data[tuple(new_keys)]
 
         # Create a new NDArray with the sliced data

From b74f471e136c897a5b4619c2943f64d55911a001 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Tue, 28 Oct 2025 18:15:46 -0400
Subject: [PATCH 62/90] fix offsets/mask shape mismatch bug in visualizing
 masked stores

---
 nki-examples/matmul.py              | 18 ++++++++++--------
 triton_viz/clients/tracer/tracer.py |  3 ++-
 2 files changed, 12 insertions(+), 9 deletions(-)

diff --git a/nki-examples/matmul.py b/nki-examples/matmul.py
index 770ea129..c4990b75 100644
--- a/nki-examples/matmul.py
+++ b/nki-examples/matmul.py
@@ -1,9 +1,11 @@
 from neuronxcc import nki
 import neuronxcc.nki.language as nl
+
 import triton_viz
 from triton_viz.clients import Tracer
 from triton_viz.core.trace import launches
 import numpy as np
+import math
 
 
 def matmul_kernel(lhs, rhs, result):
@@ -29,12 +31,12 @@ def matmul_kernel(lhs, rhs, result):
     TILE_N = 4
 
     # Use affine_range to loop over tiles
-    for m in nl.affine_range(M // TILE_M):
-        for n in nl.affine_range(N // TILE_N):
+    for m in nl.affine_range(math.ceil(M / TILE_M)):
+        for n in nl.affine_range(math.ceil(N / TILE_N)):
             # Allocate a tensor in PSUM
             res_psum = nl.zeros((TILE_M, TILE_N), nl.int32, buffer=nl.psum)
 
-            for k in nl.affine_range(K // TILE_K):
+            for k in nl.affine_range(math.ceil(K / TILE_K)):
                 # Declare the tiles on SBUF
                 lhs_tile = nl.ndarray((TILE_K, TILE_M), dtype=lhs.dtype, buffer=nl.sbuf)
                 rhs_tile = nl.ndarray((TILE_K, TILE_N), dtype=rhs.dtype, buffer=nl.sbuf)
@@ -66,21 +68,21 @@ def matmul_kernel(lhs, rhs, result):
                 mask=out_mask,
             )
 
-    return result
-
 
 TRITON_VIZ = True
 kernel_grid = (1, 1, 1)
 lhs_small = np.arange(16).astype(np.float32).reshape(4, 4)
 rhs_small = np.arange(32).astype(np.float32).reshape(4, 8)
+# lhs_small = np.arange(9).astype(np.float32).reshape(3, 3)
+# rhs_small = np.arange(18).astype(np.float32).reshape(3, 6)
 result = np.empty((lhs_small.shape[0], rhs_small.shape[1]), dtype=lhs_small.dtype)
 kernel_args = (lhs_small, rhs_small, result)
 
 if TRITON_VIZ:
     print("Executing matmul_kernel with NKI interpreter...")
     traced_kernel = triton_viz.trace(clients=Tracer(), backend="nki")(matmul_kernel)
-    kernel = traced_kernel[kernel_grid]
-    kernel(*kernel_args)
+    kernel_instance = traced_kernel[kernel_grid]
+    kernel_instance(*kernel_args)
 
     print(f"Number of launches: {len(launches)}")
     if launches:
@@ -97,7 +99,7 @@ def matmul_kernel(lhs, rhs, result):
 
     # Try to launch visualization
     try:
-        triton_viz.launch(share=True, port=8005)
+        triton_viz.launch(share=False)
     except Exception as e:
         print(f"\nError during visualization: {e}")
         import traceback
diff --git a/triton_viz/clients/tracer/tracer.py b/triton_viz/clients/tracer/tracer.py
index ba1f405d..683069ba 100644
--- a/triton_viz/clients/tracer/tracer.py
+++ b/triton_viz/clients/tracer/tracer.py
@@ -165,11 +165,12 @@ def pre_masked_store_callback(
             if not self.sample:
                 return
             keys = _convert_keys_to_numpy(keys)
-            offsets = masked_load(ptr.get_offsets().data, 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)
 
             self.records.append(Store(ptr.data_ptr(), offsets, mask_data))
 

From 1910b78077d73f326f5c827aeca29fa8ffdb281c Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Wed, 29 Oct 2025 17:38:17 -0400
Subject: [PATCH 63/90] unify DSL patching pt 1. (make nki matmul + load_store
 examples work)

---
 triton_viz/clients/tracer/tracer.py |  16 +-
 triton_viz/core/client.py           |  25 +-
 triton_viz/core/data.py             |  10 +
 triton_viz/core/nki.py              |   8 +-
 triton_viz/core/patch.py            | 367 ++++++++++++++++------------
 triton_viz/core/trace.py            |   4 +-
 6 files changed, 254 insertions(+), 176 deletions(-)

diff --git a/triton_viz/clients/tracer/tracer.py b/triton_viz/clients/tracer/tracer.py
index 683069ba..ec9cb210 100644
--- a/triton_viz/clients/tracer/tracer.py
+++ b/triton_viz/clients/tracer/tracer.py
@@ -7,6 +7,8 @@
     ReduceSum,
     Dot,
     Grid,
+    MaskedLoad,
+    MaskedStore,
     RawLoad,
     RawStore,
     Array,
@@ -238,15 +240,15 @@ def post_flip_callback(ret, x, *args, **kwargs):
             rec.call_path = _extract_user_frames()
             self.records.append(rec)
 
-        if op_type is Array:  # THTODO: only for NKI
+        if op_type is Array:
             return OpCallbacks(after_callback=post_array_callback)
-        # if op_type is Load:
-        #     return OpCallbacks(before_callback=pre_load_callback)
-        # elif op_type is Store:
-        #     return OpCallbacks(before_callback=pre_store_callback)
-        if op_type is Load:  # THTODO: only for NKI
+        elif op_type is Load:
+            return OpCallbacks(before_callback=pre_load_callback)
+        elif op_type is MaskedLoad:
             return OpCallbacks(before_callback=pre_masked_load_callback)
-        elif op_type is Store:  # THTODO: only for NKI
+        elif op_type is Store:
+            return OpCallbacks(before_callback=pre_store_callback)
+        elif op_type is MaskedStore:
             return OpCallbacks(before_callback=pre_masked_store_callback)
         elif op_type is RawLoad:
             return OpCallbacks(before_callback=pre_raw_load_callback)
diff --git a/triton_viz/core/client.py b/triton_viz/core/client.py
index 112b3fc3..d1254257 100644
--- a/triton_viz/core/client.py
+++ b/triton_viz/core/client.py
@@ -10,12 +10,11 @@
     unpatch_op,
     patch_for_loop,
     unpatch_for_loop,
-    op_list,
     patch_calls,
 )
 from functools import wraps
 from .callbacks import OpCallbacks, ForLoopCallbacks
-# from .patch import patch_lang, unpatch_lang
+from .patch import patch_lang, unpatch_lang, OPERATION_REGISTRY
 
 
 class Client(ABC):
@@ -127,27 +126,31 @@ def wrapped(*args, **kwargs):
             jit_fn.warmup = jit_fn.warmup.__wrapped__
 
     @contextmanager
-    def patch_run(self, fn):
-        with patch_calls():
+    def patch_run(self, fn, backend: str):
+        with patch_calls(backend):
             for client in self.clients.values():
-                for op in op_list:
-                    # patch ops
+                # get operations for the specified backend
+                backend_ops = OPERATION_REGISTRY[backend]["op_list"]
 
-                    callbacks = client.register_op_callback(op)
-                    patch_op(op, callbacks)
+                for op in backend_ops:
+                    # patch ops
+                    callbacks = client.register_op_callback(op, backend=backend)
+                    patch_op(op, callbacks, backend=backend)
 
                 # patch for loops
                 loop_callbacks = client.register_for_loop_callback()
                 patch_for_loop(loop_callbacks)
                 # Remaps core language functions to interpreted ones
-                # patch_lang(fn)
+                patch_lang(fn, backend)
             try:
                 yield
             finally:
-                for op in op_list:
+                backend_ops = OPERATION_REGISTRY[backend]["op_list"]
+
+                for op in backend_ops:
                     unpatch_op(op)
                 unpatch_for_loop()
-                # unpatch_lang()
+                unpatch_lang(backend)
 
     def pre_run_callback(self, fn: Callable) -> bool:
         rets = []
diff --git a/triton_viz/core/data.py b/triton_viz/core/data.py
index bdaae954..88a1689a 100644
--- a/triton_viz/core/data.py
+++ b/triton_viz/core/data.py
@@ -62,6 +62,11 @@ class Store(Op):
     masks: npt.NDArray[np.bool_]
 
 
+@dataclass
+class MaskedStore(Op):
+    name: ClassVar[str] = "masked_store"
+
+
 @dataclass
 class RawLoad(Op):
     name: ClassVar[str] = "raw_load"
@@ -75,6 +80,11 @@ class Load(Op):
     masks: npt.NDArray[np.bool_]
 
 
+@dataclass
+class MaskedLoad(Op):
+    name: ClassVar[str] = "masked_load"
+
+
 @dataclass
 class UnaryOp(Op):
     name: ClassVar[str] = "unary_op"
diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index abb2ab1c..1a46a340 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -377,7 +377,7 @@ def range(self, stop):
 nki_builder = Builder()
 
 
-def patch():
+def nki_patch_lang():
     nl.ndarray = nki_builder.ndarray
     nl.program_id = nki_builder.program_id
     nl.arange = nki_builder.arange
@@ -424,7 +424,7 @@ def patch():
     nl.device_print = print
 
 
-def unpatch():
+def nki_unpatch_lang():
     # reload the original functions
     import importlib
 
@@ -461,8 +461,6 @@ def run(self, *args, **kwargs):
         if client_manager is not None:
             client_manager.grid_callback(grid_dims)
 
-        patch()  # NKI interpreter patching
-
         # Apply AST transformer to convert nl.load/nl.store calls to nl.masked_load/nl.masked_store
         if hasattr(self.fn, "__code__"):
             # Get the source code of the function
@@ -510,5 +508,3 @@ def run(self, *args, **kwargs):
                     self.fn(*args, **kwargs)
                     if not client_manager.post_run_callback(self.fn):
                         return
-
-        unpatch()
diff --git a/triton_viz/core/patch.py b/triton_viz/core/patch.py
index b3f672a0..96a2bb81 100644
--- a/triton_viz/core/patch.py
+++ b/triton_viz/core/patch.py
@@ -2,6 +2,7 @@
 from contextlib import contextmanager
 from collections.abc import Callable
 from typing import Any, Optional
+from functools import partialmethod
 from tqdm import tqdm
 
 from .config import config as cfg
@@ -11,8 +12,10 @@
     Array,
     RawLoad,
     Load,
+    MaskedLoad,
     RawStore,
     Store,
+    MaskedStore,
     UnaryOp,
     BinaryOp,
     TernaryOp,
@@ -57,133 +60,161 @@
 from triton.runtime import JITFunction
 from triton_viz.core.nki import nki_builder
 
-BUILDER = nki_builder
-if BUILDER == interpreter_builder:
-    op_list = [
-        ProgramId,
-        RawStore,
-        Store,
-        RawLoad,
-        Load,
-        UnaryOp,
-        BinaryOp,
-        TernaryOp,
-        Dot,
-        MakeRange,
-        AddPtr,
-        Splat,
-        ExpandDims,
-        Broadcast,
-        ReduceMax,
-        ReduceMin,
-        ReduceSum,
-        MakeBlockPointer,
-        TensorPointerLoad,
-        TensorPointerStore,
-        Idiv,
-        Rsqrt,
-        CastImpl,
-        Reshape,
-        Join,
-        Fabs,
-        Ashr,
-        Advance,
-        FpToFp,
-        Umulhi,
-        Trans,
-        CumSum,
-        Bitcast,
-        AtomicCas,
-    ]
-    original_ops = {
-        ProgramId: interpreter_builder.create_get_program_id,
-        RawStore: interpreter_builder.create_store,
-        Store: interpreter_builder.create_masked_store,
-        RawLoad: interpreter_builder.create_load,
-        Load: interpreter_builder.create_masked_load,
-        Dot: interpreter_builder.create_dot,
-        UnaryOp: interpreter_builder.unary_op,
-        BinaryOp: interpreter_builder.binary_op,
-        TernaryOp: interpreter_builder.ternary_op,
-        MakeRange: interpreter_builder.create_make_range,
-        AddPtr: interpreter_builder.create_addptr,
-        ExpandDims: interpreter_builder.create_expand_dims,
-        Broadcast: interpreter_builder.create_broadcast,
-        Splat: interpreter_builder.create_splat,
-        MakeBlockPointer: interpreter_builder.create_make_block_ptr,
-        TensorPointerLoad: interpreter_builder.create_tensor_pointer_load,
-        TensorPointerStore: interpreter_builder.create_tensor_pointer_store,
-        Idiv: interpreter_builder.create_idiv,
-        Rsqrt: interpreter_builder.create_rsqrt,
-        CastImpl: interpreter_builder.cast_impl,
-        Reshape: interpreter_builder.create_reshape,
-        Join: interpreter_builder.create_join,
-        Fabs: interpreter_builder.create_fabs,
-        Ashr: interpreter_builder.create_ashr,
-        Advance: interpreter_builder.create_advance,
-        FpToFp: interpreter_builder.create_fp_to_fp,
-        Umulhi: interpreter_builder.create_umulhi,
-        Trans: interpreter_builder.create_trans,
-        Bitcast: interpreter_builder.create_bitcast,
-        AtomicCas: interpreter_builder.create_atomic_cas,
-    }
-    # Hardcoded operation attribute names to avoid issues with lambda functions
-    _OP_ATTR_NAMES = {
-        ProgramId: "create_get_program_id",
-        RawStore: "create_store",
-        Store: "create_masked_store",
-        RawLoad: "create_load",
-        Load: "create_masked_load",
-        Dot: "create_dot",
-        UnaryOp: "unary_op",
-        BinaryOp: "binary_op",
-        TernaryOp: "ternary_op",
-        MakeRange: "create_make_range",
-        AddPtr: "create_addptr",
-        ExpandDims: "create_expand_dims",
-        Broadcast: "create_broadcast",
-        Splat: "create_splat",
-        MakeBlockPointer: "create_make_block_ptr",
-        TensorPointerLoad: "create_tensor_pointer_load",
-        TensorPointerStore: "create_tensor_pointer_store",
-        Idiv: "create_idiv",
-        Rsqrt: "create_rsqrt",
-        CastImpl: "cast_impl",
-        Reshape: "create_reshape",
-        Join: "create_join",
-        Fabs: "create_fabs",
-        Ashr: "create_ashr",
-        Advance: "create_advance",
-        FpToFp: "create_fp_to_fp",
-        Umulhi: "create_umulhi",
-        Trans: "create_trans",
-        Bitcast: "create_bitcast",
-        AtomicCas: "create_atomic_cas",
-    }
-elif BUILDER == nki_builder:
-    op_list = [Array, ProgramId, Store, Load, Dot, UnaryOp, MakeRange]
-    original_ops = {
-        ProgramId: nki_builder.program_id,
-        Array: nki_builder.ndarray,
-        Load: nki_builder.masked_load,
-        Store: nki_builder.masked_store,
-        Dot: nki_builder.matmul,
-        UnaryOp: nki_builder._unary_op,
-        # BinaryOp: nki_builder.binary_op,
-        # TernaryOp: nki_builder.ternary_op,
-        MakeRange: nki_builder.arange,
-        # Rsqrt: nki_builder.create_rsqrt,
-        # CastImpl: nki_builder.cast_impl,
-    }
-    _OP_ATTR_NAMES = {
-        ProgramId: "program_id",
-        Array: "ndarray",
-        Load: "masked_load",
-        Store: "masked_store",
-        Dot: "matmul",
-        UnaryOp: "_unary_op",
-        MakeRange: "arange",
-    }
+# shared operation registry for both backends
+OPERATION_REGISTRY = {
+    "triton": {
+        "op_list": [
+            ProgramId,
+            RawStore,
+            Store,
+            RawLoad,
+            Load,
+            UnaryOp,
+            BinaryOp,
+            TernaryOp,
+            Dot,
+            MakeRange,
+            AddPtr,
+            Splat,
+            ExpandDims,
+            Broadcast,
+            ReduceMax,
+            ReduceMin,
+            ReduceSum,
+            MakeBlockPointer,
+            TensorPointerLoad,
+            TensorPointerStore,
+            Idiv,
+            Rsqrt,
+            CastImpl,
+            Reshape,
+            Join,
+            Fabs,
+            Ashr,
+            Advance,
+            FpToFp,
+            Umulhi,
+            Trans,
+            CumSum,
+            Bitcast,
+            AtomicCas,
+        ],
+        "original_ops": {
+            ProgramId: interpreter_builder.create_get_program_id,
+            RawStore: interpreter_builder.create_store,
+            Store: interpreter_builder.create_masked_store,
+            RawLoad: interpreter_builder.create_load,
+            Load: interpreter_builder.create_masked_load,
+            Dot: interpreter_builder.create_dot,
+            UnaryOp: interpreter_builder.unary_op,
+            BinaryOp: interpreter_builder.binary_op,
+            TernaryOp: interpreter_builder.ternary_op,
+            MakeRange: interpreter_builder.create_make_range,
+            AddPtr: interpreter_builder.create_addptr,
+            ExpandDims: interpreter_builder.create_expand_dims,
+            Broadcast: interpreter_builder.create_broadcast,
+            Splat: interpreter_builder.create_splat,
+            MakeBlockPointer: interpreter_builder.create_make_block_ptr,
+            TensorPointerLoad: interpreter_builder.create_tensor_pointer_load,
+            TensorPointerStore: interpreter_builder.create_tensor_pointer_store,
+            Idiv: interpreter_builder.create_idiv,
+            Rsqrt: interpreter_builder.create_rsqrt,
+            CastImpl: interpreter_builder.cast_impl,
+            Reshape: interpreter_builder.create_reshape,
+            Join: interpreter_builder.create_join,
+            Fabs: interpreter_builder.create_fabs,
+            Ashr: interpreter_builder.create_ashr,
+            Advance: interpreter_builder.create_advance,
+            FpToFp: interpreter_builder.create_fp_to_fp,
+            Umulhi: interpreter_builder.create_umulhi,
+            Trans: interpreter_builder.create_trans,
+            Bitcast: interpreter_builder.create_bitcast,
+            AtomicCas: interpreter_builder.create_atomic_cas,
+        },
+        "op_attr_names": {
+            ProgramId: "create_get_program_id",
+            RawStore: "create_store",
+            Store: "create_masked_store",
+            RawLoad: "create_load",
+            Load: "create_masked_load",
+            Dot: "create_dot",
+            UnaryOp: "unary_op",
+            BinaryOp: "binary_op",
+            TernaryOp: "ternary_op",
+            MakeRange: "create_make_range",
+            AddPtr: "create_addptr",
+            ExpandDims: "create_expand_dims",
+            Broadcast: "create_broadcast",
+            Splat: "create_splat",
+            MakeBlockPointer: "create_make_block_ptr",
+            TensorPointerLoad: "create_tensor_pointer_load",
+            TensorPointerStore: "create_tensor_pointer_store",
+            Idiv: "create_idiv",
+            Rsqrt: "create_rsqrt",
+            CastImpl: "cast_impl",
+            Reshape: "create_reshape",
+            Join: "create_join",
+            Fabs: "create_fabs",
+            Ashr: "create_ashr",
+            Advance: "create_advance",
+            FpToFp: "create_fp_to_fp",
+            Umulhi: "create_umulhi",
+            Trans: "create_trans",
+            Bitcast: "create_bitcast",
+            AtomicCas: "create_atomic_cas",
+        },
+    },
+    "nki": {
+        "op_list": [Array, ProgramId, MaskedStore, MaskedLoad, Dot, UnaryOp, MakeRange],
+        "original_ops": {
+            ProgramId: nki_builder.program_id,
+            Array: nki_builder.ndarray,
+            MaskedLoad: nki_builder.masked_load,
+            MaskedStore: nki_builder.masked_store,
+            Dot: nki_builder.matmul,
+            UnaryOp: nki_builder._unary_op,
+            MakeRange: nki_builder.arange,
+        },
+        "op_attr_names": {
+            ProgramId: "program_id",
+            Array: "ndarray",
+            MaskedLoad: "masked_load",
+            MaskedStore: "masked_store",
+            Dot: "matmul",
+            UnaryOp: "_unary_op",
+            MakeRange: "arange",
+        },
+    },
+}
+
+BUILDER = interpreter_builder
+current_backend = "triton"
+op_list = OPERATION_REGISTRY[current_backend]["op_list"]
+original_ops = OPERATION_REGISTRY[current_backend]["original_ops"]  # type: ignore
+_OP_ATTR_NAMES = OPERATION_REGISTRY[current_backend]["op_attr_names"]  # type: ignore
+
+
+def get_builder_for_backend(backend: str):
+    """Get the appropriate builder for a given backend."""
+    return nki_builder if backend == "nki" else interpreter_builder
+
+
+def detect_current_backend():
+    """Detect the current backend based on which operations have been patched."""
+    # Check if Triton operations have been patched first (since this example uses Triton)
+    if hasattr(interpreter_builder, "create_get_program_id") and hasattr(
+        interpreter_builder.create_get_program_id, "__wrapped__"
+    ):
+        return "triton"
+    # Check if NKI operations have been patched
+    elif hasattr(nki_builder, "program_id") and hasattr(
+        nki_builder.program_id, "__wrapped__"
+    ):
+        return "nki"
+    # Default to checking BUILDER global
+    else:
+        return "nki" if BUILDER == nki_builder else "triton"
+
 
 reduce_map: dict[type[Op], Callable] = {
     ReduceMax: tl.max,
@@ -237,24 +268,34 @@ def __call__(self, *args, **kwargs):
         return ret
 
 
-def patch_op(op_type: type[Op], callbacks: OpCallbacks):
+def patch_op(op_type: type[Op], callbacks: OpCallbacks, backend: Optional[str] = None):
     """
     Register a callback to be called before and after an operator is executed.
 
     :param op_type: The type of the operator to register the callback for.
     :param callbacks: The OpCallbacks object containing before_callback, after_callback, and op_overrider.
+    :param backend: The backend to use ('triton', 'nki', or None for current backend).
     """
-    if op_type in original_ops:
-        # create a new function that calls the before_callback, the original op and the after_callback
-        op_name = _OP_ATTR_NAMES[op_type]
-        original_op = original_ops[op_type]
+    if backend is None:
+        backend = current_backend
+
+    if backend not in OPERATION_REGISTRY:
+        raise ValueError(f"Unknown backend: {backend}")
+
+    backend_ops = OPERATION_REGISTRY[backend]["original_ops"]  # type: ignore
+    backend_attr_names = OPERATION_REGISTRY[backend]["op_attr_names"]  # type: ignore
+    backend_builder = nki_builder if backend == "nki" else interpreter_builder
+
+    if op_type in backend_ops:
+        op_name = backend_attr_names[op_type]  # type: ignore
+        original_op = backend_ops[op_type]  # type: ignore
         patched_op = PatchOp(original_op, op_type, callbacks)
         setattr(
-            BUILDER,
+            backend_builder,
             op_name,
             lambda *args, **kwargs: patched_op(*args, **kwargs),
         )
-    elif op_type in reduce_map or op_type in scan_map:
+    elif backend == "triton" and (op_type in reduce_map or op_type in scan_map):
         if op_type in reduce_map:
             op_name = reduce_map[op_type].__name__
         elif op_type in scan_map:
@@ -262,7 +303,7 @@ def patch_op(op_type: type[Op], callbacks: OpCallbacks):
         original_op = getattr(tl, op_name)
         patched_op = PatchOp(original_op, op_type, callbacks)
         setattr(tl, op_name, lambda *args, **kwargs: patched_op(*args, **kwargs))
-    elif op_type in math_map:
+    elif backend == "triton" and op_type in math_map:
         op_name = math_map[op_type].__name__
         original_op = getattr(tl.math, op_name)
         patched_op = PatchOp(original_op, op_type, callbacks)
@@ -278,9 +319,9 @@ def unpatch_op(op_type: type[Op]):
     :param op_type: The type of the operator to unregister the callback for.
     """
     if op_type in original_ops:
-        original_op = original_ops[op_type]
+        original_op = original_ops[op_type]  # type: ignore
         # Use hardcoded name from _OP_ATTR_NAMES
-        op_name = _OP_ATTR_NAMES[op_type]
+        op_name = _OP_ATTR_NAMES[op_type]  # type: ignore
         setattr(BUILDER, op_name, original_op)
 
 
@@ -453,8 +494,18 @@ def unpatch_for_loop():
     _loop_patcher.unpatch()
 
 
-def patch_lang(fn):
-    triton_patch_lang(fn)
+def patch_lang(fn, backend):
+    if backend == "triton":
+        triton_patch_lang(fn)
+    elif backend == "nki":
+        from triton_viz.core.nki import nki_patch_lang
+
+        nki_patch_lang()
+    else:
+        raise ValueError(
+            f"Unsupported backend {backend} received. Triton-viz only supports one of ('triton', 'nki')."
+        )
+
     fn.__globals__["_triton_viz_loop_patcher"] = _loop_patcher
     # Wrap tl.flip to emit a Flip record after computing result
     try:
@@ -547,15 +598,21 @@ def _maybe_list(arr):
         pass
 
 
-def unpatch_lang():
+def unpatch_lang(backend):
     import importlib
     import sys
 
-    if tl.__name__ in sys.modules:
-        importlib.reload(tl)
+    if backend == "triton":
+        if tl.__name__ in sys.modules:
+            importlib.reload(tl)
+    elif backend == "nki":
+        from triton_viz.core.nki import nki_unpatch_lang
 
+        nki_unpatch_lang()
 
-def _grid_executor_call(self, *args_dev, **kwargs):
+
+def _grid_executor_call(self, *args_dev, backend=None, **kwargs):
+    assert backend is not None
     if kwargs.pop("warmup", False):
         return
 
@@ -578,7 +635,7 @@ def run_grid_loops(grid):
                     leave=False,
                     disable=not (cfg.report_grid_execution_progress and grid[2] > 1),
                 ):
-                    BUILDER.set_grid_idx(x, y, z)
+                    builder.set_grid_idx(x, y, z)
                     client_manager.grid_idx_callback((x, y, z))
                     if not client_manager.pre_run_callback(self.fn):
                         return
@@ -617,24 +674,34 @@ def run_grid_loops(grid):
     grid = self.grid(call_args) if callable(self.grid) else self.grid
     assert len(grid) <= 3
     grid = grid + (1,) * (3 - len(grid))
-    BUILDER.set_grid_dim(*grid)
+
+    # Use the correct builder based on current backend
+    builder = get_builder_for_backend(backend)
+
+    if backend == "nki":
+        builder.set_grid_dim(grid)
+    else:  # triton
+        builder.set_grid_dim(*grid)
     client_manager.grid_callback(grid)
     run_grid_loops(grid)
     # Copy arguments back to propagate side-effects
     self._restore_args_dev(args_dev, args_hst, kwargs, kwargs_hst)
 
 
-def _jit_function_call(self, *args, **kwargs):
-    patch_lang(self.fn)
+def _jit_function_call(
+    self, *args, backend=None, **kwargs
+):  # NOTE: is this ever called?
+    assert backend is not None
+    patch_lang(self.fn, backend)
     return self.fn(*args, **kwargs)
 
 
 @contextmanager
-def patch_calls():
+def patch_calls(backend):
     old_grid_executor_call = GridExecutor.__call__
     old_jit_function_call = JITFunction.__call__
-    GridExecutor.__call__ = _grid_executor_call
-    JITFunction.__call__ = _jit_function_call
+    GridExecutor.__call__ = partialmethod(_grid_executor_call, backend=backend)
+    JITFunction.__call__ = partialmethod(_jit_function_call, backend=backend)
     try:
         yield
     finally:
diff --git a/triton_viz/core/trace.py b/triton_viz/core/trace.py
index 1053c3d8..d0e78550 100644
--- a/triton_viz/core/trace.py
+++ b/triton_viz/core/trace.py
@@ -105,7 +105,7 @@ def run(self, *args, **kwargs):
             if self.warmup_runner:
                 self.warmup_runner.warmup(*args, **kwargs)
 
-        with self.client_manager.patch_run(self.base_fn):
+        with self.client_manager.patch_run(self.base_fn, backend="triton"):
             kwargs.update({"client_manager": self.client_manager})
             kwargs.update({"jit_fn": self.jit_fn})
             ret = self.runner.run(*args, **kwargs)
@@ -145,7 +145,7 @@ def __call__(self, *args, **kwargs):
         return self[(1, 1, 1)](*args, **kwargs)
 
     def run(self, *args, **kwargs):
-        with self.client_manager.patch_run(self.interpreter_fn):
+        with self.client_manager.patch_run(self.func, backend="nki"):
             kwargs.update({"client_manager": self.client_manager})
             ret = self.interpreter_fn.run(*args, **kwargs)
             self.finalize()

From a120447f8be6eae4617a5ff2f6376c603eda1801 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Fri, 31 Oct 2025 13:51:16 -0400
Subject: [PATCH 64/90] remove some files

---
 {nki-examples => examples/nki}/matmul.py |   0
 nki-examples/rmsnorm.py                  |  94 ---
 nki-examples/softmax.py                  |  71 ---
 nki-examples/test_kernels.py             | 230 -------
 tests/nki/test_masked_load.py            | 770 +++++++++++++++++++++++
 tests/{ => nki}/test_nki.py              |   0
 tests/test_masked_load.py                | 585 -----------------
 7 files changed, 770 insertions(+), 980 deletions(-)
 rename {nki-examples => examples/nki}/matmul.py (100%)
 delete mode 100644 nki-examples/rmsnorm.py
 delete mode 100644 nki-examples/softmax.py
 delete mode 100644 nki-examples/test_kernels.py
 create mode 100644 tests/nki/test_masked_load.py
 rename tests/{ => nki}/test_nki.py (100%)
 delete mode 100644 tests/test_masked_load.py

diff --git a/nki-examples/matmul.py b/examples/nki/matmul.py
similarity index 100%
rename from nki-examples/matmul.py
rename to examples/nki/matmul.py
diff --git a/nki-examples/rmsnorm.py b/nki-examples/rmsnorm.py
deleted file mode 100644
index 854d6304..00000000
--- a/nki-examples/rmsnorm.py
+++ /dev/null
@@ -1,94 +0,0 @@
-from torch_xla.core import xla_model as xm
-import neuronxcc.nki as nki
-import neuronxcc.nki.language as nl
-import numpy as np
-import math
-import torch
-
-@nki.jit
-def nki_rmsnorm_kernel(a_tensor, g_tensor):
-  # Calculate out_tensor = a_tensor/RMS(a_tensor) * g_tensor
-  # Where RMS(a_tensor) = sqrt((1/N) * sum(a_tensor * a_tensor))
-  # and N = a_tensor.shape[1]
-  # Reduction (mean) is performed in the free (2nd) dimension
-  B, D = a_tensor.shape
-  out_tensor = nl.ndarray(a_tensor.shape, dtype=a_tensor.dtype,
-                          buffer=nl.shared_hbm)
-
-  # Make sure shapes match
-  assert D == g_tensor.shape[0]
-
-  # Generate tensor indices to index input tensor
-  ix = nl.arange(2)[:, None]
-  iw = nl.arange(1)[:, None]
-  #iy = nl.arange(a_tensor.shape[1])[None, :]
-  iy = nl.arange(8)[None, :]
-
-  # Load RMSNorm weight once, reused by rows/tiles of a_tensor
-  g_tile = nl.load(g_tensor.reshape((1, D))[iw, iy], mask=iy < D)
-
-  # Process 2 rows at a time due to 2-partition tile size limitation
-  # Since we're not reducing across the first dimension
-  # Tiles can be processed independently
-  for i in nl.affine_range(math.ceil(B/2)):
-
-    # Load input data from external memory to on-chip memory
-    mask = (i * 2 + ix < B) & (iy < D)
-    a_tile = nl.load(a_tensor[i * 2 + ix, iy],
-                    mask=mask)
-
-    # Compute element-wise square of a_tensor
-    in_square = nl.square(a_tile)
-
-    # Calculate sum of squared elements, along last dimension
-    square_sum = nl.sum(in_square, axis=[1], mask=mask)
-
-    # Scale and get a reciprocal
-    mean = square_sum / D
-
-    # Take square root of mean and then reciprocal with
-    # rsqrt API (one ISA instruction)
-    rms_reciprocal = nl.rsqrt(mean)
-
-    # Scale the input tensor
-    out_tile = nl.multiply(a_tile, rms_reciprocal)
-
-    # Broadcast weight along first axis to match tensor shape
-    # B_active = min(B - i * 2, 2)
-    g_bcast = g_tile.broadcast_to((2, 8))
-
-    # Multiply with the RMSNorm weight
-    out_tile[...] = nl.multiply(out_tile, g_bcast,
-                           mask=(i * 2 + ix < B))
-
-    # store the addition results back to external memory (out_tensor)
-    nl.store(out_tensor[i * 2 + ix, iy], value=out_tile,
-            mask=mask)
-
-  return out_tensor
-
-# ref
-def torch_rmsnorm_kernel(a_tensor, g_tensor):
-  # Square the tensor (element-wise)
-  in_square = a_tensor.pow(2)
-  # Calculate means in the free dimension
-  mean = in_square.mean(dim=1, keepdim=True)
-  # Scale by reciprocal of sqrt(mean)
-  tensor = a_tensor * torch.rsqrt(mean)
-
-  # Scale the output by the weight
-  return tensor * g_tensor
-
-device = 'cpu'
-
-a_tensor = torch.arange(15).float().view(3, 5).to(device=device)
-g_tensor = torch.arange(5).float().to(device=device)
-
-output_nki = nki.simulate_kernel(nki_rmsnorm_kernel, a_tensor.numpy(), g_tensor.numpy())
-
-output_torch = torch_rmsnorm_kernel(a_tensor, g_tensor).numpy()
-
-if np.allclose(output_torch, output_nki, atol=1e-5, rtol=1e-3):
-  print("NKI and Torch match")
-else:
-  print("NKI and Torch differ")
diff --git a/nki-examples/softmax.py b/nki-examples/softmax.py
deleted file mode 100644
index 628ee4b6..00000000
--- a/nki-examples/softmax.py
+++ /dev/null
@@ -1,71 +0,0 @@
-from neuronxcc import nki
-import neuronxcc.nki.language as nl
-
-import torch
-import triton_viz
-import torch_xla.core.xla_model as xm
-from triton_viz.clients import Tracer
-from triton_viz.core import config as cfg
-from triton_viz.core.trace import launches
-import numpy as np
-import math
-
-
-@nki.jit
-def softmax(in_tensor):
-    # assume B,D; softmax on D for now
-    B, D = in_tensor.shape
-    out_tensor = nl.ndarray((B, D), dtype=in_tensor.dtype, buffer=nl.shared_hbm)
-
-    #assert nl.tile_size.pmax == 128
-    num_tiles = math.ceil(B / 128)
-    for tile_idx in nl.affine_range(num_tiles):
-        i_p = tile_idx * 128 + nl.arange(128)[:, None]
-        i_f = nl.arange(D)[None, :]
-        mask = i_p < B
-        tile = nl.load(in_tensor[i_p, i_f], mask=mask)
-        tile_exp = nl.exp(tile)
-        tile_expsum = nl.sum(tile_exp, -1, keepdims=True, mask=mask)
-        tile_softmax = tile_exp / tile_expsum
-        nl.store(out_tensor[i_p, i_f], value=tile_softmax, mask=mask)
-    return out_tensor
-
-
-if __name__ == "__main__":
-    cfg.reset()
-    device = "cpu"
-    B, D = 16, 32
-    torch.manual_seed(0)
-    x = torch.rand((B, D), dtype=torch.float32, device=device)
-    y = np.exp(x) / np.exp(x).sum(-1, keepdim=True)
-    TRITON_VIZ = False
-
-    if TRITON_VIZ:
-        softmax = triton_viz.trace()(softmax)
-        softmax[(1,1,1)](x)
-
-        # Print records to see what's being captured
-        print(f"Number of launches: {len(launches)}")
-        if launches:
-            launch = launches[-1]
-            print(f"Number of records: {len(launch.records)}")
-            for i, record in enumerate(launch.records):
-                print(f"Record {i}: {type(record).__name__}")
-                if hasattr(record, "ptr"):
-                    print(f"  ptr: {record.ptr}")
-                if hasattr(record, "offsets"):
-                    print(f"  offsets shape: {record.offsets.shape}")
-                if hasattr(record, "masks"):
-                    print(f"  masks shape: {record.masks.shape}")
-
-        # Try to launch visualization
-        try:
-            triton_viz.launch()
-        except Exception as e:
-            print(f"\nError during visualization: {e}")
-            import traceback
-
-            traceback.print_exc()
-    else:
-        y_pred = nki.simulate_kernel(softmax, x.numpy())
-        assert np.allclose(y, y_pred)
diff --git a/nki-examples/test_kernels.py b/nki-examples/test_kernels.py
deleted file mode 100644
index 48f439db..00000000
--- a/nki-examples/test_kernels.py
+++ /dev/null
@@ -1,230 +0,0 @@
-import math
-import neuronxcc.nki.language as nl
-import neuronxcc.nki as nki
-import numpy as np
-import torch
-from triton_viz.clients import Tracer
-import triton_viz
-from triton_viz.core.trace import launches
-
-
-def add_kernel(a, b):  # fails @ ix < B
-    pid_x = nl.program_id(0)
-    pid_y = nl.program_id(1)
-    B, D = a.shape
-
-    # memory shared across all SPMD instances
-    c_output = nl.ndarray(a.shape, buffer=nl.shared_hbm, dtype=a.dtype)
-
-    ix = pid_x * 128 + nl.arange(128)[:, None]
-    iy = pid_y * 512 + nl.arange(512)[None, :]
-    mask = (ix < B) & (iy < D)
-    a_tmp = nl.load(a[ix, iy], mask=mask)
-    b_tmp = nl.load(b[ix, iy], mask=mask)
-    c_tmp = a_tmp + b_tmp
-    nl.store(c_output[ix, iy], value=c_tmp, mask=mask)
-    return c_output
-
-def copy_kernel(a):  # fails @ ix < B
-    B, D = a.shape
-    out = nl.ndarray(a.shape, dtype=a.dtype, buffer=nl.shared_hbm)
-
-    pid_x = nl.program_id(0)
-    pid_y = nl.program_id(1)
-
-    ix = pid_x * 128 + nl.arange(128)[:, None]
-    iy = pid_y * 512 + nl.arange(512)[None, :]
-    mask = (ix < B) & (iy < D)
-
-    a_tmp = nl.load(a[ix, iy], mask=mask)
-    nl.store(out[ix, iy], value=a_tmp, mask=mask)
-    return out
-
-def print_kernel():  # works
-    a = nl.ndarray([4, 4], dtype=nl.float32, buffer=nl.shared_hbm)
-    y = nl.ndarray(
-        [4, 4],
-        dtype=np.float32,
-    )
-    nl.store(a, value=y)
-    print(a)
-    return a
-
-def tmp0_kernel(a):  # works
-    B, D = a.shape
-    out = nl.ndarray(a.shape, dtype=a.dtype, buffer=nl.shared_hbm)
-
-    pid_x = nl.program_id(0)
-    pid_y = nl.program_id(1)
-    nl.device_print(f"pid_x:", pid_x);
-    nl.device_print(f"pid_y:", pid_y);
-
-    ix = pid_x * 2 + nl.arange(2)[:, None]
-    iy = pid_y * 4 + nl.arange(4)[None, :]
-    mask = (ix < B) & (iy < D)
-    a_tmp = nl.load(a[ix, iy], mask=mask)
-    nl.device_print(f"a_tmp:", a_tmp)
-
-    ix2 = pid_x * 3 + nl.arange(3)[:, None]
-    iy2 = pid_y * 4 + nl.arange(4)[None, :]
-    mask2 = (ix2 < B) & (iy2 < D)
-    a_tmp2 = nl.load(a[ix2, iy2], mask=mask2)
-    nl.device_print(f"a_tmp2:", a_tmp2)
-
-    iy3 = nl.arange(3)[None, :] < -1
-    a_tmp3 = nl.load(a[80, nl.arange(3)[None, :]], mask=iy3)
-    nl.device_print(f"a_tmp3:", a_tmp3)
-
-    nl.store(out[ix, iy], value=a_tmp, mask=mask)
-    # for load(src, mask), src.shape, mask.shape need to be same, return shape
-    # for store(dst, value, mask), dst.shape, value.shape, mask.shape need to be same, return shape
-    return out
-
-def mgrid_kernel(a):  # test nl.mgrid functionality
-    B, D = a.shape
-    out = nl.ndarray(a.shape, dtype=a.dtype, buffer=nl.shared_hbm)
-
-    coords = nl.mgrid[:B, :D]
-    mask = (coords[0] < B) & (coords[1] < D)
-    a_tmp = nl.load(a[coords[0], coords[1]], mask=mask)
-    nl.store(out[coords[0], coords[1]], value=2*a_tmp, mask=mask)
-
-    return out
-
-def xyz_kernel(a):  # works
-    B, T, C, H, W = a.shape
-    out = nl.ndarray(a.shape, dtype=a.dtype, buffer=nl.shared_hbm)
-
-    pid_x = nl.program_id(0)
-    pid_y = nl.program_id(1)
-    pid_z = nl.program_id(2)
-
-    for i in range(H // 128):
-        for j in range(W // 512):
-            i_h = pid_x * 128 + nl.arange(128)[:, None]
-            i_n = pid_y * 512 + nl.arange(512)[None, :]
-            mask = (i_h < H) & (i_n < W)
-            a_tmp = nl.load(a[pid_x, pid_y, pid_z, i_h, i_n], mask=mask)
-            nl.store(out[pid_x, pid_y, pid_z, i_h, i_n], value=a_tmp, mask=mask)
-    return out
-
-def tmp1_kernel(a):
-    B, D = a.shape
-    out = nl.ndarray(a.shape, dtype=a.dtype, buffer=nl.shared_hbm)
-
-    pid_x = nl.program_id(0)
-    pid_y = nl.program_id(1)
-
-    ix = pid_x * 2 + nl.arange(2)[:, None]
-    iy = pid_y * 4 + nl.arange(4)[None, :]
-    mask = (ix < B) & (iy < D)
-    a_tmp = nl.load(a[ix, iy], mask=mask)
-
-    #iy3 = nl.arange(3)[None, :] < -1
-    #iy3 = nl.mgrid[80:83, 80:85] < -1
-    iy3 = (nl.arange(3)[:, None] < -1) & (nl.arange(5)[None, :] < -1)
-    a_tmp3 = nl.load(a[80:83, 80:85], mask=iy3)
-    nl.device_print(f"a_tmp3:", a_tmp3)
-
-    #nl.store(out[ix, iy], value=a_tmp, mask=mask)
-    return out
-
-
-B, D = 1024, 1024
-x = torch.rand((B, D))
-y = torch.rand((B, D))
-
-blocks_x = math.ceil(B / 128)
-blocks_y = math.ceil(D / 512)
-
-kernel = tmp1_kernel
-if kernel == add_kernel:
-    kernel_grid = (blocks_x, blocks_y)
-    kernel_args = (x.numpy(), y.numpy())
-    z1 = x + y
-if kernel == copy_kernel:
-    kernel_grid = (blocks_x, blocks_y)
-    kernel_args = (x.numpy(),)
-    z1 = x
-elif kernel == print_kernel:
-    kernel_grid = (1, 1, 1)
-    kernel_args = ()
-    z1 = x
-if kernel == tmp0_kernel:
-    #B, D = 129, 512
-    B, D = 3, 5
-    #x = torch.rand((B, D))
-    #y = torch.rand((B, D))
-    x = torch.arange(B*D).int().reshape(B, D)
-    y = -torch.arange(B*D).int().reshape(B, D)
-
-    #blocks_x = math.ceil(B / 128)
-    blocks_x = math.ceil(B / 2)
-    blocks_y = math.ceil(D / 4)
-
-    kernel_grid = (blocks_x, blocks_y)
-    kernel_args = (x.numpy(),)
-    z1 = x
-if kernel == xyz_kernel:
-    B, T, C, H, W = 2, 3, 4, 1024, 1024
-    x = torch.rand((B, T, C, H, W))
-    kernel_grid = (B, T, C)
-    kernel_args = (x.numpy(),)
-    z1 = x
-if kernel == tmp1_kernel:
-    B, D = 3, 5
-    x = torch.arange(B*D).int().reshape(B, D)
-    y = -torch.arange(B*D).int().reshape(B, D)
-
-    blocks_x = math.ceil(B / 2)
-    blocks_y = math.ceil(D / 4)
-
-    kernel_grid = (blocks_x, blocks_y)
-    kernel_args = (x.numpy(),)
-    z1 = x
-if kernel == mgrid_kernel:
-    B, D = 1024, 1024
-    x = torch.rand((B, D))
-
-    # Expected result: simple copy (nl.mgrid is just tested with device_print)
-    z1 = 2 * x  # Should be identical to input
-
-    kernel_grid = (1, 1, 1)
-    kernel_args = (x.numpy(),)
-
-TRITON_VIZ = True
-
-if TRITON_VIZ:
-    kernel = triton_viz.trace(clients=Tracer(), backend="nki")(kernel)
-    kk = kernel[kernel_grid]
-    z2 = kk(*kernel_args)
-    #print((z1 - z2).abs().max())
-
-    #print(f"Number of launches: {len(launches)}")
-    #if launches:
-    #    launch = launches[-1]
-    #    print(f"Number of records: {len(launch.records)}")
-    #    for i, record in enumerate(launch.records):
-    #        print(f"Record {i}: {type(record).__name__}")
-    #        if hasattr(record, "ptr"):
-    #            print(f"  ptr: {record.ptr}")
-    #        if hasattr(record, "offsets"):
-    #            print(f"  offsets shape: {record.offsets.shape}")
-    #        if hasattr(record, "masks"):
-    #            print(f"  masks shape: {record.masks.shape}")
-
-    ## Try to launch visualization
-    #try:
-    #    triton_viz.launch(share=False)
-    #except Exception as e:
-    #    print(f"\nError during visualization: {e}")
-    #    import traceback
-
-    #    traceback.print_exc()
-else:
-    kernel = nki.jit(kernel)
-    z2 = nki.simulate_kernel(kernel[kernel_grid], *kernel_args)
-
-z2 = torch.from_numpy(z2)
-print((z1 - z2).abs().max())
-assert torch.allclose(z1, z2)
diff --git a/tests/nki/test_masked_load.py b/tests/nki/test_masked_load.py
new file mode 100644
index 00000000..1c0ebf3b
--- /dev/null
+++ b/tests/nki/test_masked_load.py
@@ -0,0 +1,770 @@
+import numpy as np
+import pytest
+from triton_viz.core.nki_masked_load import masked_load, masked_store
+
+
+def print_op_details(
+    test_name,
+    operation_type,
+    input_data,
+    keys,
+    values=None,
+    mask=None,
+    output=None,
+    error_expected=None,
+):
+    """Print detailed information about a masked load/store operation."""
+    print(f"{test_name}:")
+    print(f"Operation type: {operation_type}")
+    print("---------------------------------")
+    print(f"Input:\n{input_data}")
+    print("---------------------------------")
+    print(f"Values:\n{values}")
+    print("---------------------------------")
+    print(f"Keys:\n{keys}")
+    print("---------------------------------")
+    print(f"Mask:\n{mask}")
+    print("---------------------------------")
+    print(f"Result {operation_type}:\n{output}")
+    if error_expected:
+        print(f"Expecting {error_expected}...")
+    print()
+
+
+def test_masked_load_none_case():
+    """Test masked_load with mask=None (direct indexing)."""
+    arr = np.array([1, 2, 3, 4, 5])
+    result = masked_load(arr, (slice(1, 4),), mask=None)
+    expected = arr[1:4]
+
+    print_op_details(
+        "Test 1: mask=None case", "load", arr, (slice(1, 4),), mask=None, output=result
+    )
+
+    assert np.array_equal(result, expected)
+
+
+def test_masked_load_in_bounds():
+    """Test masked_load with in-bounds indexing and mask."""
+    arr = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
+    mask = np.array([[True, True, False], [True, False, True], [False, True, True]])
+    result = masked_load(arr, (slice(0, 3), slice(0, 3)), mask=mask)
+    result[
+        [0, 1, 2], [2, 1, 0]
+    ] = 0  # manually set masked out positions to 0 for testing
+    expected = np.array([[1, 2, 0], [4, 0, 6], [0, 8, 9]])
+
+    print_op_details(
+        "Test 2: In-bounds indexing with mask",
+        "load",
+        arr,
+        (slice(0, 3), slice(0, 3)),
+        mask=mask,
+        output=result,
+    )
+
+    assert np.array_equal(result, expected)
+
+
+def test_masked_load_out_of_bounds():
+    """Test masked_load with out-of-bounds indexing and mask."""
+    arr = np.array([[1, 2], [3, 4]])
+    mask = np.array(
+        [
+            [True, False, False],
+            [False, True, False],
+            [False, False, False],
+            [False, False, False],
+            [False, False, False],
+        ]
+    )  # note that arr.shape != mask.shape, this is intended
+
+    # normally this arr[slice] would OOB but the mask at the OOB idxs are false so they're not loaded
+    result = masked_load(arr, (slice(0, 5), slice(0, 3)), mask=mask)
+
+    print_op_details(
+        "Test 3: Out-of-bounds indexing with mask",
+        "load",
+        arr,
+        (slice(0, 5), slice(0, 3)),
+        mask=mask,
+        output=result,
+    )
+
+    assert result.shape == (5, 3)
+
+
+def test_masked_load_integer_oob():
+    """Test masked_load with integer out-of-bounds indexing."""
+    arr = np.array([10, 20, 30])
+    mask = np.array([False])
+    result = masked_load(arr, ([5],), mask=mask)  # idx 5 is OOB
+
+    print_op_details(
+        "Test 4: Integer indexing out of bounds",
+        "load",
+        arr,
+        ([5],),
+        mask=mask,
+        output=result,
+    )
+
+    assert isinstance(result, np.ndarray)
+
+
+def test_masked_load_array_indexing():
+    """Test masked_load with array indexing."""
+    arr = np.array([100, 200, 300, 400, 500])
+    indices = np.array([0, 2, 4])
+    mask = np.array([True, True, True])
+    result = masked_load(arr, (indices,), mask=mask)
+    expected = arr[indices]
+
+    print_op_details(
+        "Test 5: Array indexing", "load", arr, (indices,), mask=mask, output=result
+    )
+
+    assert np.array_equal(result, expected)
+
+
+def test_masked_load_mixed_indexing():
+    """Test masked_load with mixed indexing types."""
+    arr = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])
+    mask = np.array([True, False, True, False])
+    result = masked_load(arr, (slice(0, 4), 2), mask=mask)  # set col 2
+
+    print_op_details(
+        "Test 6: Mixed indexing types",
+        "load",
+        arr,
+        (slice(0, 4), 2),
+        mask=mask,
+        output=result,
+    )
+
+    assert result.shape == (4,)
+
+
+def test_masked_load_complex_indexing():
+    """Test masked_load with complex multi-dimensional indexing."""
+    arr = np.array(
+        [
+            [
+                [1, 2, 3],
+                [4, 5, 6],
+                [7, 8, 9],
+            ],
+            [
+                [11, 12, 13],
+                [14, 15, 16],
+                [17, 18, 19],
+            ],
+            [
+                [21, 22, 23],
+                [24, 25, 26],
+                [27, 28, 29],
+            ],
+        ]
+    )
+    arr_slice = (
+        slice(None, None, None),
+        None,
+        np.arange(20)[:, None],
+        np.arange(3)[None, :],
+        None,
+        None,
+    )
+    mask = np.mgrid[:3, :1, :20, :3, :1, :1]
+    mask = (mask[0] > 0) & (mask[2] < 3) & (mask[3] < 2)
+    result = masked_load(arr, arr_slice, mask=mask)
+    UD = np.iinfo(arr.dtype).max  # UD = undefined (out of bounds)
+    expected = np.array(
+        [
+            [
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+            ],
+            [
+                [11, 12, UD],
+                [14, 15, UD],
+                [17, 18, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+            ],
+            [
+                [21, 22, UD],
+                [24, 25, UD],
+                [27, 28, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+            ],
+        ]
+    )[:, None, :, :, None, None]
+
+    print(
+        "Load Test 7: Complex multi-dimensional indexing (omitted details, check code to see inputs/outputs)"
+    )
+
+    assert np.allclose(result, expected)
+
+
+def test_masked_load_incomplete_slices():
+    """Test masked_load with incomplete slices."""
+    arr = np.array(
+        [
+            [
+                [1, 2, 3],
+                [4, 5, 6],
+                [7, 8, 9],
+            ],
+            [
+                [11, 12, 13],
+                [14, 15, 16],
+                [17, 18, 19],
+            ],
+            [
+                [21, 22, 23],
+                [24, 25, 26],
+                [27, 28, 29],
+            ],
+        ]
+    )
+    arr_slice = (slice(None, None, None),)
+    mask = np.mgrid[:3, :3, :3]
+    mask = (0 < mask[0]) & (mask[0] < 2) & (mask[1] < 2) & (mask[2] != 1)
+    result = masked_load(arr, arr_slice, mask=mask)
+    UD = np.iinfo(arr.dtype).max  # UD = undefined (out of bounds)
+    expected = np.array(
+        [
+            [
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+            ],
+            [
+                [11, UD, 13],
+                [14, UD, 16],
+                [UD, UD, UD],
+            ],
+            [
+                [UD, UD, UD],
+                [UD, UD, UD],
+                [UD, UD, UD],
+            ],
+        ]
+    )
+
+    print_op_details(
+        "Test 8: Incomplete slices", "load", arr, arr_slice, mask=mask, output=result
+    )
+
+    assert np.allclose(result, expected)
+
+
+def test_masked_load_index_error_with_true_mask():
+    """Test that IndexError is raised when mask=True at out-of-bounds location."""
+    arr = np.array(
+        [
+            [
+                [1, 2, 3],
+                [4, 5, 6],
+                [7, 8, 9],
+            ],
+            [
+                [11, 12, 13],
+                [14, 15, 16],
+                [17, 18, 19],
+            ],
+            [
+                [21, 22, 23],
+                [24, 25, 26],
+                [27, 28, 29],
+            ],
+        ]
+    )
+    arr_slice = (slice(None, None, None), slice(0, 4))
+    mask = np.mgrid[:3, :4, :3]
+    mask = (mask[0] < 3) & (mask[1] < 3) & (mask[2] < 3)
+    mask[-1, -1, -1] = True  # index error
+
+    print_op_details(
+        "Test 9: Index error should still happen when mask=True at OOB location",
+        "load",
+        arr,
+        arr_slice,
+        mask=mask,
+        error_expected="IndexError",
+    )
+
+    with pytest.raises(IndexError):
+        masked_load(arr, arr_slice, mask=mask)
+    print("Success: Correctly raised IndexError\n")
+
+
+def test_masked_load_mask_shape_mismatch():
+    """Test that AssertionError is raised when mask shape doesn't match array slice."""
+    arr = np.array(
+        [
+            [
+                [1, 2, 3],
+                [4, 5, 6],
+                [7, 8, 9],
+            ],
+            [
+                [11, 12, 13],
+                [14, 15, 16],
+                [17, 18, 19],
+            ],
+            [
+                [21, 22, 23],
+                [24, 25, 26],
+                [27, 28, 29],
+            ],
+        ]
+    )
+    arr_slice = (slice(None, None, None),)
+    mask = np.mgrid[:3, :4, :3]
+    mask = (mask[0] < 3) & (mask[1] < 3) & (mask[2] < 3)
+
+    print_op_details(
+        "Test 10: Error if mask shape wrong",
+        "load",
+        arr,
+        arr_slice,
+        mask=mask,
+        error_expected="AssertionError",
+    )
+
+    with pytest.raises(AssertionError):
+        masked_load(arr, arr_slice, mask=mask)
+    print("Success: Correctly raised AssertionError\n")
+
+
+def test_masked_load_correct_mask_shape():
+    """Test masked_load with correctly shaped mask."""
+    arr = np.array(
+        [
+            [
+                [1, 2, 3],
+                [4, 5, 6],
+                [7, 8, 9],
+            ],
+            [
+                [11, 12, 13],
+                [14, 15, 16],
+                [17, 18, 19],
+            ],
+            [
+                [21, 22, 23],
+                [24, 25, 26],
+                [27, 28, 29],
+            ],
+        ]
+    )
+    arr_slice = (slice(None, None, None),)
+    mask = np.arange(27).reshape(3, 3, 3) % 2 == 0
+    result = masked_load(arr, arr_slice, mask=mask)
+
+    print_op_details(
+        "Test 11: Correct mask shape", "load", arr, arr_slice, mask=mask, output=result
+    )
+
+    assert result.shape == mask.shape
+
+
+def test_masked_store_none_case():
+    """Test masked_store with mask=None (direct indexing)."""
+    arr = np.array([1, 2, 3, 4, 5])
+    values = np.array([10, 20, 30])
+    arr_copy = arr.copy()
+    masked_store(arr_copy, (slice(1, 4),), values, mask=None)
+    expected = arr.copy()
+    expected[1:4] = values
+
+    print_op_details(
+        "Store Test 1: mask=None case",
+        "store",
+        arr,
+        (slice(1, 4),),
+        values=values,
+        mask=None,
+        output=arr_copy,
+    )
+
+    assert np.array_equal(arr_copy, expected)
+
+
+def test_masked_store_in_bounds():
+    """Test masked_store with in-bounds indexing and mask."""
+    arr = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
+    values = np.array([[100, 200, 0], [400, 0, 600], [0, 800, 900]])
+    mask = np.array([[True, True, False], [True, False, True], [False, True, True]])
+    arr_copy = arr.copy()
+    masked_store(arr_copy, (slice(0, 3), slice(0, 3)), values, mask=mask)
+    expected = np.array([[100, 200, 3], [400, 5, 600], [7, 800, 900]])
+
+    print_op_details(
+        "Store Test 2: In-bounds indexing with mask",
+        "store",
+        arr,
+        (slice(0, 3), slice(0, 3)),
+        values=values,
+        mask=mask,
+        output=arr_copy,
+    )
+
+    assert np.array_equal(arr_copy, expected)
+
+
+def test_masked_store_out_of_bounds():
+    """Test masked_store with out-of-bounds indexing and mask."""
+    arr = np.array([[1, 2], [3, 4]])
+    values = np.array(
+        [
+            [100, 0, 0],
+            [0, 200, 0],
+            [0, 0, 0],
+            [0, 0, 0],
+            [0, 0, 0],
+        ]
+    )
+    mask = np.array(
+        [
+            [True, False, False],
+            [False, True, False],
+            [False, False, False],
+            [False, False, False],
+            [False, False, False],
+        ]
+    )
+    arr_copy = arr.copy()
+    masked_store(arr_copy, (slice(0, 5), slice(0, 3)), values, mask=mask)
+    expected = np.array([[100, 2], [3, 200]])
+
+    print_op_details(
+        "Store Test 3: Out-of-bounds indexing with mask",
+        "store",
+        arr,
+        (slice(0, 5), slice(0, 3)),
+        values=values,
+        mask=mask,
+        output=arr_copy,
+    )
+
+    assert np.array_equal(arr_copy, expected)
+
+
+def test_masked_store_integer_oob():
+    """Test masked_store with integer out-of-bounds indexing."""
+    arr = np.array([10, 20, 30])
+    values = np.array([999])
+    mask = np.array([False])  # mask=False at OOB index
+    arr_copy = arr.copy()
+    masked_store(arr_copy, ([5],), values, mask=mask)  # Index 5 is OOB
+
+    print_op_details(
+        "Store Test 4: Integer indexing out of bounds",
+        "store",
+        arr,
+        ([5],),
+        values=values,
+        mask=mask,
+        output=arr_copy,
+    )
+
+    assert np.array_equal(arr_copy, arr)
+
+
+def test_masked_store_array_indexing():
+    """Test masked_store with array indexing."""
+    arr = np.array([100, 200, 300, 400, 500])
+    indices = np.array([0, 2, 4])
+    values = np.array([111, 333, 555])
+    mask = np.array([True, True, True])
+    arr_copy = arr.copy()
+    masked_store(arr_copy, (indices,), values, mask=mask)
+    expected = np.array([111, 200, 333, 400, 555])
+
+    print_op_details(
+        "Store Test 5: Array indexing",
+        "store",
+        arr,
+        (indices,),
+        values=values,
+        mask=mask,
+        output=arr_copy,
+    )
+
+    assert np.array_equal(arr_copy, expected)
+
+
+def test_masked_store_mixed_indexing():
+    """Test masked_store with mixed indexing types."""
+    arr = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])
+    values = np.array([100, 200, 300, 400])
+    mask = np.array([True, False, True, False])
+    arr_copy = arr.copy()
+    masked_store(arr_copy, (slice(0, 4), 2), values, mask=mask)  # set col 2
+    expected = np.array([[1, 2, 100, 4], [5, 6, 7, 8], [9, 10, 300, 12]])
+
+    print_op_details(
+        "Store Test 6: Mixed indexing types",
+        "store",
+        arr,
+        (slice(0, 4), 2),
+        values=values,
+        mask=mask,
+        output=arr_copy,
+    )
+
+    assert np.array_equal(arr_copy, expected)
+
+
+def test_masked_store_complex_indexing():
+    """Test masked_store with complex multi-dimensional indexing."""
+    arr = np.array(
+        [
+            [
+                [1, 2, 3],
+                [4, 5, 6],
+                [7, 8, 9],
+            ],
+            [
+                [11, 12, 13],
+                [14, 15, 16],
+                [17, 18, 19],
+            ],
+            [
+                [21, 22, 23],
+                [24, 25, 26],
+                [27, 28, 29],
+            ],
+        ]
+    )
+    arr_slice = (
+        slice(None, None, None),
+        None,
+        np.arange(20)[:, None],
+        np.arange(3)[None, :],
+        None,
+        None,
+    )
+    mask = np.mgrid[:3, :1, :20, :3, :1, :1]
+    mask = (mask[0] > 0) & (mask[2] < 3) & (mask[3] < 2)
+    values = np.array(
+        [
+            [
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+            ],
+            [
+                [31, 32, 99],
+                [34, 35, 99],
+                [37, 38, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+            ],
+            [
+                [41, 42, 99],
+                [44, 45, 99],
+                [47, 48, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+                [99, 99, 99],
+            ],
+        ]
+    )[:, None, :, :, None, None]
+    arr_copy = arr.copy()
+    masked_store(arr_copy, arr_slice, values, mask=mask)
+    expected_modified = np.array(
+        [
+            [
+                [1, 2, 3],
+                [4, 5, 6],
+                [7, 8, 9],
+            ],
+            [
+                [31, 32, 13],
+                [34, 35, 16],
+                [37, 38, 19],
+            ],
+            [
+                [41, 42, 23],
+                [44, 45, 26],
+                [47, 48, 29],
+            ],
+        ]
+    )
+
+    print(
+        "Store Test 7: Complex multi-dimensional indexing (omitted details, check code to see inputs/outputs)"
+    )
+
+    assert np.array_equal(arr_copy, expected_modified)
+
+
+def test_masked_store_oob_with_true_mask():
+    """Test that IndexError is raised when mask=True at out-of-bounds location."""
+    arr = np.array([[1, 2], [3, 4]])
+    values = np.array([[100, 200], [300, 400], [500, 600]])
+    mask = np.array(
+        [
+            [True, True],
+            [True, True],
+            [True, False],  # This position is OOB and mask=True
+        ]
+    )
+    arr_copy = arr.copy()
+
+    print_op_details(
+        "Store Test 8: OOB with mask=True should raise IndexError",
+        "store",
+        arr,
+        (slice(0, 3), slice(0, 2)),
+        values=values,
+        mask=mask,
+        error_expected="IndexError",
+    )
+
+    with pytest.raises(IndexError):
+        masked_store(arr_copy, (slice(0, 3), slice(0, 2)), values, mask=mask)
+    print("Success: Correctly raised IndexError\n")
+
+
+def test_masked_store_values_mask_shape_mismatch():
+    """Test that AssertionError is raised when values and mask shapes don't match."""
+    arr = np.array([1, 2, 3])
+    values = np.array([10, 20, 30])  # Shape (3,)
+    mask = np.array([True, False])  # Shape (2,) - different from values
+    arr_copy = arr.copy()
+
+    print_op_details(
+        "Store Test 9: Values and mask shape mismatch should raise AssertionError",
+        "store",
+        arr,
+        (slice(0, 3),),
+        values=values,
+        mask=mask,
+        error_expected="AssertionError",
+    )
+
+    with pytest.raises(AssertionError):
+        masked_store(arr_copy, (slice(0, 3),), values, mask=mask)
+    print("Success: Correctly raised AssertionError\n")
+
+
+def test_masked_store_values_shape_mismatch():
+    """Test that IndexError is raised when values shape is incompatible."""
+    arr = np.array([1, 2, 3])
+    values = np.array([10, 20, 30, 40])  # Wrong shape
+    mask = np.array([True, False, True, False])
+    arr_copy = arr.copy()
+
+    print_op_details(
+        "Store Test 10: Values shape mismatch",
+        "store",
+        arr,
+        (slice(0, 3),),
+        values=values,
+        mask=mask,
+        error_expected="IndexError",
+    )
+
+    with pytest.raises(IndexError):
+        masked_store(arr_copy, (slice(0, 3),), values, mask=mask)
+    print("Success: Correctly raised IndexError\n")
diff --git a/tests/test_nki.py b/tests/nki/test_nki.py
similarity index 100%
rename from tests/test_nki.py
rename to tests/nki/test_nki.py
diff --git a/tests/test_masked_load.py b/tests/test_masked_load.py
deleted file mode 100644
index 5eedc3cf..00000000
--- a/tests/test_masked_load.py
+++ /dev/null
@@ -1,585 +0,0 @@
-import numpy as np
-import pytest
-from triton_viz.core.nki_masked_load import masked_load, masked_store
-
-def print_op_details(test_name, operation_type, input_data, keys, values=None, mask=None, output=None, error_expected=None):
-    """Print detailed information about a masked load/store operation."""
-    print(f"{test_name}:")
-    print(f"Operation type: {operation_type}")
-    print("---------------------------------")
-    print(f"Input:\n{input_data}")
-    print("---------------------------------")
-    print(f"Values:\n{values}")
-    print("---------------------------------")
-    print(f"Keys:\n{keys}")
-    print("---------------------------------")
-    print(f"Mask:\n{mask}")
-    print("---------------------------------")
-    print(f"Result {operation_type}:\n{output}")
-    if error_expected:
-        print(f"Expecting {error_expected}...")
-    print()
-
-def test_masked_load_none_case():
-    """Test masked_load with mask=None (direct indexing)."""
-    arr = np.array([1, 2, 3, 4, 5])
-    result = masked_load(arr, (slice(1, 4),), mask=None)
-    expected = arr[1:4]
-
-    print_op_details("Test 1: mask=None case", "load", arr, (slice(1, 4),), mask=None, output=result)
-
-    assert np.array_equal(result, expected)
-
-def test_masked_load_in_bounds():
-    """Test masked_load with in-bounds indexing and mask."""
-    arr = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
-    mask = np.array([[True, True, False], [True, False, True], [False, True, True]])
-    result = masked_load(arr, (slice(0, 3), slice(0, 3)), mask=mask)
-    result[[0,1,2],[2,1,0]] = 0  # manually set masked out positions to 0 for testing
-    expected = np.array([[1, 2, 0], [4, 0, 6], [0, 8, 9]])
-
-    print_op_details("Test 2: In-bounds indexing with mask", "load", arr, (slice(0, 3), slice(0, 3)), mask=mask, output=result)
-
-    assert np.array_equal(result, expected)
-
-def test_masked_load_out_of_bounds():
-    """Test masked_load with out-of-bounds indexing and mask."""
-    arr = np.array([[1, 2], [3, 4]])
-    mask = np.array([
-        [True , False, False],
-        [False, True , False],
-        [False, False, False],
-        [False, False, False],
-        [False, False, False],
-    ])  # note that arr.shape != mask.shape, this is intended
-
-    # normally this arr[slice] would OOB but the mask at the OOB idxs are false so they're not loaded
-    result = masked_load(arr, (slice(0, 5), slice(0, 3)), mask=mask)
-
-    print_op_details("Test 3: Out-of-bounds indexing with mask", "load", arr, (slice(0, 5), slice(0, 3)), mask=mask, output=result)
-
-    assert result.shape == (5, 3)
-
-def test_masked_load_integer_oob():
-    """Test masked_load with integer out-of-bounds indexing."""
-    arr = np.array([10, 20, 30])
-    mask = np.array([False])
-    result = masked_load(arr, ([5],), mask=mask)  # idx 5 is OOB
-
-    print_op_details("Test 4: Integer indexing out of bounds", "load", arr, ([5],), mask=mask, output=result)
-
-    assert isinstance(result, np.ndarray)
-
-def test_masked_load_array_indexing():
-    """Test masked_load with array indexing."""
-    arr = np.array([100, 200, 300, 400, 500])
-    indices = np.array([0, 2, 4])
-    mask = np.array([True, True, True])
-    result = masked_load(arr, (indices,), mask=mask)
-    expected = arr[indices]
-
-    print_op_details("Test 5: Array indexing", "load", arr, (indices,), mask=mask, output=result)
-
-    assert np.array_equal(result, expected)
-
-def test_masked_load_mixed_indexing():
-    """Test masked_load with mixed indexing types."""
-    arr = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])
-    mask = np.array([True, False, True, False])
-    result = masked_load(arr, (slice(0, 4), 2), mask=mask)  # set col 2
-
-    print_op_details("Test 6: Mixed indexing types", "load", arr, (slice(0, 4), 2), mask=mask, output=result)
-
-    assert result.shape == (4,)
-
-def test_masked_load_complex_indexing():
-    """Test masked_load with complex multi-dimensional indexing."""
-    arr = np.array([
-        [
-            [1, 2, 3],
-            [4, 5, 6],
-            [7, 8, 9],
-        ],
-        [
-            [11, 12, 13],
-            [14, 15, 16],
-            [17, 18, 19],
-        ],
-        [
-            [21, 22, 23],
-            [24, 25, 26],
-            [27, 28, 29],
-        ],
-    ])
-    arr_slice = (
-        slice(None, None, None),
-        None,
-        np.arange(20)[:, None],
-        np.arange(3)[None, :],
-        None,
-        None
-    )
-    mask = np.mgrid[:3, :1, :20, :3, :1, :1]
-    mask = (mask[0] > 0) & (mask[2] < 3) & (mask[3] < 2)
-    result = masked_load(arr, arr_slice, mask=mask)
-    UD = np.iinfo(arr.dtype).max  # UD = undefined (out of bounds)
-    expected = np.array([
-        [
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-        ],
-        [
-            [11, 12, UD],
-            [14, 15, UD],
-            [17, 18, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-        ],
-        [
-            [21, 22, UD],
-            [24, 25, UD],
-            [27, 28, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-            ],
-        ]
-    )[:, None, :, :, None, None]
-
-    print("Load Test 7: Complex multi-dimensional indexing (omitted details, check code to see inputs/outputs)")
-
-    assert np.allclose(result, expected)
-
-def test_masked_load_incomplete_slices():
-    """Test masked_load with incomplete slices."""
-    arr = np.array([
-        [
-            [1, 2, 3],
-            [4, 5, 6],
-            [7, 8, 9],
-        ],
-        [
-            [11, 12, 13],
-            [14, 15, 16],
-            [17, 18, 19],
-        ],
-        [
-            [21, 22, 23],
-            [24, 25, 26],
-            [27, 28, 29],
-        ],
-    ])
-    arr_slice = (slice(None, None, None),)
-    mask = np.mgrid[:3, :3, :3]
-    mask = (0 < mask[0]) & (mask[0] < 2) & (mask[1] < 2) & (mask[2] != 1)
-    result = masked_load(arr, arr_slice, mask=mask)
-    UD = np.iinfo(arr.dtype).max  # UD = undefined (out of bounds)
-    expected = np.array([
-        [
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-        ],
-        [
-            [11, UD, 13],
-            [14, UD, 16],
-            [UD, UD, UD],
-        ],
-        [
-            [UD, UD, UD],
-            [UD, UD, UD],
-            [UD, UD, UD],
-        ],
-    ])
-
-    print_op_details("Test 8: Incomplete slices", "load", arr, arr_slice, mask=mask, output=result)
-
-    assert np.allclose(result, expected)
-
-def test_masked_load_index_error_with_true_mask():
-    """Test that IndexError is raised when mask=True at out-of-bounds location."""
-    arr = np.array([
-        [
-            [1, 2, 3],
-            [4, 5, 6],
-            [7, 8, 9],
-        ],
-        [
-            [11, 12, 13],
-            [14, 15, 16],
-            [17, 18, 19],
-        ],
-        [
-            [21, 22, 23],
-            [24, 25, 26],
-            [27, 28, 29],
-        ],
-    ])
-    arr_slice = (slice(None, None, None), slice(0, 4))
-    mask = np.mgrid[:3, :4, :3]
-    mask = (mask[0] < 3) & (mask[1] < 3) & (mask[2] < 3)
-    mask[-1, -1, -1] = True  # index error
-
-    print_op_details("Test 9: Index error should still happen when mask=True at OOB location", "load", arr, arr_slice, mask=mask, error_expected="IndexError")
-
-    with pytest.raises(IndexError):
-        masked_load(arr, arr_slice, mask=mask)
-    print("Success: Correctly raised IndexError\n")
-
-def test_masked_load_mask_shape_mismatch():
-    """Test that AssertionError is raised when mask shape doesn't match array slice."""
-    arr = np.array([
-        [
-            [1, 2, 3],
-            [4, 5, 6],
-            [7, 8, 9],
-        ],
-        [
-            [11, 12, 13],
-            [14, 15, 16],
-            [17, 18, 19],
-        ],
-        [
-            [21, 22, 23],
-            [24, 25, 26],
-            [27, 28, 29],
-        ],
-    ])
-    arr_slice = (slice(None, None, None),)
-    mask = np.mgrid[:3, :4, :3]
-    mask = (mask[0] < 3) & (mask[1] < 3) & (mask[2] < 3)
-
-    print_op_details("Test 10: Error if mask shape wrong", "load", arr, arr_slice, mask=mask, error_expected="AssertionError")
-
-    with pytest.raises(AssertionError):
-        masked_load(arr, arr_slice, mask=mask)
-    print("Success: Correctly raised AssertionError\n")
-
-def test_masked_load_correct_mask_shape():
-    """Test masked_load with correctly shaped mask."""
-    arr = np.array([
-        [
-            [1, 2, 3],
-            [4, 5, 6],
-            [7, 8, 9],
-        ],
-        [
-            [11, 12, 13],
-            [14, 15, 16],
-            [17, 18, 19],
-        ],
-        [
-            [21, 22, 23],
-            [24, 25, 26],
-            [27, 28, 29],
-        ],
-    ])
-    arr_slice = (slice(None, None, None),)
-    mask = np.arange(27).reshape(3, 3, 3) % 2 == 0
-    result = masked_load(arr, arr_slice, mask=mask)
-
-    print_op_details("Test 11: Correct mask shape", "load", arr, arr_slice, mask=mask, output=result)
-
-    assert result.shape == mask.shape
-
-def test_masked_store_none_case():
-    """Test masked_store with mask=None (direct indexing)."""
-    arr = np.array([1, 2, 3, 4, 5])
-    values = np.array([10, 20, 30])
-    arr_copy = arr.copy()
-    masked_store(arr_copy, (slice(1, 4),), values, mask=None)
-    expected = arr.copy()
-    expected[1:4] = values
-
-    print_op_details("Store Test 1: mask=None case", "store", arr, (slice(1, 4),), values=values, mask=None, output=arr_copy)
-
-    assert np.array_equal(arr_copy, expected)
-
-def test_masked_store_in_bounds():
-    """Test masked_store with in-bounds indexing and mask."""
-    arr = np.array([[1, 2, 3], [4, 5, 6], [7, 8, 9]])
-    values = np.array([[100, 200, 0], [400, 0, 600], [0, 800, 900]])
-    mask = np.array([[True, True, False], [True, False, True], [False, True, True]])
-    arr_copy = arr.copy()
-    masked_store(arr_copy, (slice(0, 3), slice(0, 3)), values, mask=mask)
-    expected = np.array([[100, 200, 3], [400, 5, 600], [7, 800, 900]])
-
-    print_op_details("Store Test 2: In-bounds indexing with mask", "store", arr, (slice(0, 3), slice(0, 3)), values=values, mask=mask, output=arr_copy)
-
-    assert np.array_equal(arr_copy, expected)
-
-def test_masked_store_out_of_bounds():
-    """Test masked_store with out-of-bounds indexing and mask."""
-    arr = np.array([[1, 2], [3, 4]])
-    values = np.array([
-        [100, 0, 0],
-        [0, 200, 0],
-        [0, 0, 0],
-        [0, 0, 0],
-        [0, 0, 0],
-    ])
-    mask = np.array([
-        [True , False, False],
-        [False, True , False],
-        [False, False, False],
-        [False, False, False],
-        [False, False, False],
-    ])
-    arr_copy = arr.copy()
-    masked_store(arr_copy, (slice(0, 5), slice(0, 3)), values, mask=mask)
-    expected = np.array([[100, 2], [3, 200]])
-
-    print_op_details("Store Test 3: Out-of-bounds indexing with mask", "store", arr, (slice(0, 5), slice(0, 3)), values=values, mask=mask, output=arr_copy)
-
-    assert np.array_equal(arr_copy, expected)
-
-def test_masked_store_integer_oob():
-    """Test masked_store with integer out-of-bounds indexing."""
-    arr = np.array([10, 20, 30])
-    values = np.array([999])
-    mask = np.array([False])  # mask=False at OOB index
-    arr_copy = arr.copy()
-    masked_store(arr_copy, ([5],), values, mask=mask)  # Index 5 is OOB
-
-    print_op_details("Store Test 4: Integer indexing out of bounds", "store", arr, ([5],), values=values, mask=mask, output=arr_copy)
-
-    assert np.array_equal(arr_copy, arr)
-
-def test_masked_store_array_indexing():
-    """Test masked_store with array indexing."""
-    arr = np.array([100, 200, 300, 400, 500])
-    indices = np.array([0, 2, 4])
-    values = np.array([111, 333, 555])
-    mask = np.array([True, True, True])
-    arr_copy = arr.copy()
-    masked_store(arr_copy, (indices,), values, mask=mask)
-    expected = np.array([111, 200, 333, 400, 555])
-
-    print_op_details("Store Test 5: Array indexing", "store", arr, (indices,), values=values, mask=mask, output=arr_copy)
-
-    assert np.array_equal(arr_copy, expected)
-
-def test_masked_store_mixed_indexing():
-    """Test masked_store with mixed indexing types."""
-    arr = np.array([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])
-    values = np.array([100, 200, 300, 400])
-    mask = np.array([True, False, True, False])
-    arr_copy = arr.copy()
-    masked_store(arr_copy, (slice(0, 4), 2), values, mask=mask)  # set col 2
-    expected = np.array([[1, 2, 100, 4], [5, 6, 7, 8], [9, 10, 300, 12]])
-
-    print_op_details("Store Test 6: Mixed indexing types", "store", arr, (slice(0, 4), 2), values=values, mask=mask, output=arr_copy)
-
-    assert np.array_equal(arr_copy, expected)
-
-def test_masked_store_complex_indexing():
-    """Test masked_store with complex multi-dimensional indexing."""
-    arr = np.array([
-        [
-            [1, 2, 3],
-            [4, 5, 6],
-            [7, 8, 9],
-        ],
-        [
-            [11, 12, 13],
-            [14, 15, 16],
-            [17, 18, 19],
-        ],
-        [
-            [21, 22, 23],
-            [24, 25, 26],
-            [27, 28, 29],
-        ],
-    ])
-    arr_slice = (
-        slice(None, None, None),
-        None,
-        np.arange(20)[:, None],
-        np.arange(3)[None, :],
-        None,
-        None
-    )
-    mask = np.mgrid[:3, :1, :20, :3, :1, :1]
-    mask = (mask[0] > 0) & (mask[2] < 3) & (mask[3] < 2)
-    values = np.array([
-        [
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-        ],
-        [
-            [31, 32, 99],
-            [34, 35, 99],
-            [37, 38, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-        ],
-        [
-            [41, 42, 99],
-            [44, 45, 99],
-            [47, 48, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            [99, 99, 99],
-            ],
-        ]
-    )[:, None, :, :, None, None]
-    arr_copy = arr.copy()
-    masked_store(arr_copy, arr_slice, values, mask=mask)
-    expected_modified = np.array([
-        [
-            [1, 2, 3],
-            [4, 5, 6],
-            [7, 8, 9],
-        ],
-        [
-            [31, 32, 13],
-            [34, 35, 16],
-            [37, 38, 19],
-        ],
-        [
-            [41, 42, 23],
-            [44, 45, 26],
-            [47, 48, 29],
-        ],
-    ])
-
-    print("Store Test 7: Complex multi-dimensional indexing (omitted details, check code to see inputs/outputs)")
-
-    assert np.array_equal(arr_copy, expected_modified)
-
-def test_masked_store_oob_with_true_mask():
-    """Test that IndexError is raised when mask=True at out-of-bounds location."""
-    arr = np.array([[1, 2], [3, 4]])
-    values = np.array([[100, 200], [300, 400], [500, 600]])
-    mask = np.array([
-        [True, True],
-        [True, True],
-        [True, False],  # This position is OOB and mask=True
-    ])
-    arr_copy = arr.copy()
-
-    print_op_details("Store Test 8: OOB with mask=True should raise IndexError", "store", arr, (slice(0, 3), slice(0, 2)), values=values, mask=mask, error_expected="IndexError")
-
-    with pytest.raises(IndexError):
-        masked_store(arr_copy, (slice(0, 3), slice(0, 2)), values, mask=mask)
-    print("Success: Correctly raised IndexError\n")
-
-def test_masked_store_values_mask_shape_mismatch():
-    """Test that AssertionError is raised when values and mask shapes don't match."""
-    arr = np.array([1, 2, 3])
-    values = np.array([10, 20, 30])  # Shape (3,)
-    mask = np.array([True, False])   # Shape (2,) - different from values
-    arr_copy = arr.copy()
-
-    print_op_details("Store Test 9: Values and mask shape mismatch should raise AssertionError", "store", arr, (slice(0, 3),), values=values, mask=mask, error_expected="AssertionError")
-
-    with pytest.raises(AssertionError):
-        masked_store(arr_copy, (slice(0, 3),), values, mask=mask)
-    print("Success: Correctly raised AssertionError\n")
-
-def test_masked_store_values_shape_mismatch():
-    """Test that IndexError is raised when values shape is incompatible."""
-    arr = np.array([1, 2, 3])
-    values = np.array([10, 20, 30, 40])  # Wrong shape
-    mask = np.array([True, False, True, False])
-    arr_copy = arr.copy()
-
-    print_op_details("Store Test 10: Values shape mismatch", "store", arr, (slice(0, 3),), values=values, mask=mask, error_expected="IndexError")
-
-    with pytest.raises(IndexError):
-        masked_store(arr_copy, (slice(0, 3),), values, mask=mask)
-    print("Success: Correctly raised IndexError\n")

From b5797ef9c7042d46b51dced2a1ec9914be6475ae Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Fri, 31 Oct 2025 14:00:00 -0400
Subject: [PATCH 65/90] normal args and kwargs

---
 triton_viz/clients/tracer/tracer.py | 18 +++++++-----------
 1 file changed, 7 insertions(+), 11 deletions(-)

diff --git a/triton_viz/clients/tracer/tracer.py b/triton_viz/clients/tracer/tracer.py
index ec9cb210..1ca2a7cf 100644
--- a/triton_viz/clients/tracer/tracer.py
+++ b/triton_viz/clients/tracer/tracer.py
@@ -115,7 +115,7 @@ def _extract_user_frames() -> list[traceback.FrameSummary]:
                     return [f]
             return stack[-1:]
 
-        def pre_load_callback(ptr, mask, *ignore_args, **ignore_kwargs):
+        def pre_load_callback(ptr, mask, *args, **kwargs):
             if not self.sample:
                 return
             first_ptr = np.reshape(ptr.data, (-1))[0]
@@ -133,13 +133,11 @@ def _convert_keys_to_numpy(keys):
             else:
                 return keys
 
-        def post_array_callback(ret, *ignore_args, **ignore_kwargs):
+        def post_array_callback(ret, *args, **kwargs):
             assert hasattr(ret, "data")
             self.tensors.append(ret)
 
-        def pre_masked_load_callback(
-            ptr, keys, mask=None, *ignore_args, **ignore_kwargs
-        ):
+        def pre_masked_load_callback(ptr, keys, mask=None, *args, **kwargs):
             if not self.sample:
                 return
             keys = _convert_keys_to_numpy(keys)
@@ -152,7 +150,7 @@ def pre_masked_load_callback(
                 )
             )
 
-        def pre_store_callback(ptr, value, mask, *ignore_args, **ignore_kwargs):
+        def pre_store_callback(ptr, value, mask, *args, **kwargs):
             if not self.sample:
                 return
             first_ptr = np.reshape(ptr.data, (-1))[0]
@@ -161,9 +159,7 @@ def pre_store_callback(ptr, value, mask, *ignore_args, **ignore_kwargs):
             rec.call_path = _extract_user_frames()
             self.records.append(rec)
 
-        def pre_masked_store_callback(
-            ptr, keys, value, mask=None, *ignore_args, **ignore_kwargs
-        ):
+        def pre_masked_store_callback(ptr, keys, value, mask=None, *args, **kwargs):
             if not self.sample:
                 return
             keys = _convert_keys_to_numpy(keys)
@@ -200,7 +196,7 @@ def pre_raw_store_callback(ptr, value):
             self.records.append(rec)
 
         def post_reduce_sum_callback(
-            ret, input, axis=None, keep_dims=False, *ignore_args, **ignore_kwargs
+            ret, input, axis=None, keep_dims=False, *args, **kwargs
         ):
             if not self.sample:
                 return
@@ -208,7 +204,7 @@ def post_reduce_sum_callback(
             output_shape = ret.handle.data.shape
             self.records.append(ReduceSum(input_shape, axis, keep_dims, output_shape))
 
-        def post_dot_callback(ret, input, other, *ignore_args, **ignore_kwargs):
+        def post_dot_callback(ret, input, other, *args, **kwargs):
             if not self.sample:
                 return
             input_shape = input.data.shape

From 4e0f02a0eb1a2f192f6a0bd27df3210b72c864a1 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Fri, 31 Oct 2025 18:54:19 -0400
Subject: [PATCH 66/90] allocate instead of array

---
 examples/nki/matmul.py              |  2 +-
 triton_viz/clients/tracer/tracer.py |  8 ++++----
 triton_viz/core/data.py             | 15 +++++++++++++--
 triton_viz/core/patch.py            | 16 ++++++++++++----
 4 files changed, 30 insertions(+), 11 deletions(-)

diff --git a/examples/nki/matmul.py b/examples/nki/matmul.py
index c4990b75..293c1768 100644
--- a/examples/nki/matmul.py
+++ b/examples/nki/matmul.py
@@ -107,7 +107,7 @@ def matmul_kernel(lhs, rhs, result):
         traceback.print_exc()
 else:
     print("Executing NKI JIT-ed matmul_kernel...")
-    compiled_kernel = nki.jit(matmul_kernel)
+    compiled_kernel = nki.jit(matmul_kernel, kernel_return=False)
     z2 = nki.simulate_kernel(compiled_kernel[kernel_grid], *kernel_args)
 
 z2 = result
diff --git a/triton_viz/clients/tracer/tracer.py b/triton_viz/clients/tracer/tracer.py
index 1ca2a7cf..6684b000 100644
--- a/triton_viz/clients/tracer/tracer.py
+++ b/triton_viz/clients/tracer/tracer.py
@@ -11,7 +11,7 @@
     MaskedStore,
     RawLoad,
     RawStore,
-    Array,
+    Allocate,
     Flip,
 )
 from triton_viz.core.nki_masked_load import masked_load
@@ -133,7 +133,7 @@ def _convert_keys_to_numpy(keys):
             else:
                 return keys
 
-        def post_array_callback(ret, *args, **kwargs):
+        def post_allocate_callback(ret, *args, **kwargs):
             assert hasattr(ret, "data")
             self.tensors.append(ret)
 
@@ -236,8 +236,8 @@ def post_flip_callback(ret, x, *args, **kwargs):
             rec.call_path = _extract_user_frames()
             self.records.append(rec)
 
-        if op_type is Array:
-            return OpCallbacks(after_callback=post_array_callback)
+        if op_type is Allocate:
+            return OpCallbacks(after_callback=post_allocate_callback)
         elif op_type is Load:
             return OpCallbacks(before_callback=pre_load_callback)
         elif op_type is MaskedLoad:
diff --git a/triton_viz/core/data.py b/triton_viz/core/data.py
index 88a1689a..f278501b 100644
--- a/triton_viz/core/data.py
+++ b/triton_viz/core/data.py
@@ -44,8 +44,8 @@ class ProgramId(Op):
 
 
 @dataclass
-class Array(Op):
-    name: ClassVar[str] = "array"
+class Allocate(Op):
+    name: ClassVar[str] = "allocate"
     ptr: int
 
 
@@ -60,6 +60,11 @@ class Store(Op):
     ptr: int
     offsets: npt.NDArray[np.int_]
     masks: npt.NDArray[np.bool_]
+    mem_src: str = "SBUF"
+    mem_dst: str = "HBM"
+    backend: str = "nki"
+    bytes: int = 0
+    time_idx: int = 0
 
 
 @dataclass
@@ -78,6 +83,12 @@ class Load(Op):
     ptr: int
     offsets: npt.NDArray[np.int_]
     masks: npt.NDArray[np.bool_]
+    # buffer: str
+    mem_src: str = "HBM"
+    mem_dst: str = "SBUF"
+    backend: str = "nki"
+    bytes: int = 0
+    time_idx: int = 0
 
 
 @dataclass
diff --git a/triton_viz/core/patch.py b/triton_viz/core/patch.py
index 96a2bb81..b97d00f9 100644
--- a/triton_viz/core/patch.py
+++ b/triton_viz/core/patch.py
@@ -9,7 +9,7 @@
 from .callbacks import OpCallbacks, ForLoopCallbacks
 from .data import (
     Op,
-    Array,
+    Allocate,
     RawLoad,
     Load,
     MaskedLoad,
@@ -165,10 +165,18 @@
         },
     },
     "nki": {
-        "op_list": [Array, ProgramId, MaskedStore, MaskedLoad, Dot, UnaryOp, MakeRange],
+        "op_list": [
+            Allocate,
+            ProgramId,
+            MaskedStore,
+            MaskedLoad,
+            Dot,
+            UnaryOp,
+            MakeRange,
+        ],
         "original_ops": {
             ProgramId: nki_builder.program_id,
-            Array: nki_builder.ndarray,
+            Allocate: nki_builder.ndarray,
             MaskedLoad: nki_builder.masked_load,
             MaskedStore: nki_builder.masked_store,
             Dot: nki_builder.matmul,
@@ -177,7 +185,7 @@
         },
         "op_attr_names": {
             ProgramId: "program_id",
-            Array: "ndarray",
+            Allocate: "ndarray",
             MaskedLoad: "masked_load",
             MaskedStore: "masked_store",
             Dot: "matmul",

From a913fdacc91181cd8801b7b142a77e29d2a0ceee Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Mon, 3 Nov 2025 14:29:25 -0500
Subject: [PATCH 67/90] reorder callbacks

---
 triton_viz/clients/tracer/tracer.py | 26 +++++++++++++-------------
 1 file changed, 13 insertions(+), 13 deletions(-)

diff --git a/triton_viz/clients/tracer/tracer.py b/triton_viz/clients/tracer/tracer.py
index 6684b000..5df8b588 100644
--- a/triton_viz/clients/tracer/tracer.py
+++ b/triton_viz/clients/tracer/tracer.py
@@ -115,6 +115,10 @@ def _extract_user_frames() -> list[traceback.FrameSummary]:
                     return [f]
             return stack[-1:]
 
+        def post_allocate_callback(ret, *args, **kwargs):
+            assert hasattr(ret, "data")
+            self.tensors.append(ret)
+
         def pre_load_callback(ptr, mask, *args, **kwargs):
             if not self.sample:
                 return
@@ -124,6 +128,15 @@ def pre_load_callback(ptr, mask, *args, **kwargs):
             rec.call_path = _extract_user_frames()
             self.records.append(rec)
 
+        def pre_store_callback(ptr, value, mask, *args, **kwargs):
+            if not self.sample:
+                return
+            first_ptr = np.reshape(ptr.data, (-1))[0]
+            tensor = self._get_tensor(first_ptr)
+            rec = Store(tensor.data_ptr(), ptr.data - tensor.data_ptr(), mask.data)
+            rec.call_path = _extract_user_frames()
+            self.records.append(rec)
+
         def _convert_keys_to_numpy(keys):
             """Convert any NDArrays in keys to numpy arrays."""
             if isinstance(keys, (tuple, list)):
@@ -133,10 +146,6 @@ def _convert_keys_to_numpy(keys):
             else:
                 return keys
 
-        def post_allocate_callback(ret, *args, **kwargs):
-            assert hasattr(ret, "data")
-            self.tensors.append(ret)
-
         def pre_masked_load_callback(ptr, keys, mask=None, *args, **kwargs):
             if not self.sample:
                 return
@@ -150,15 +159,6 @@ def pre_masked_load_callback(ptr, keys, mask=None, *args, **kwargs):
                 )
             )
 
-        def pre_store_callback(ptr, value, mask, *args, **kwargs):
-            if not self.sample:
-                return
-            first_ptr = np.reshape(ptr.data, (-1))[0]
-            tensor = self._get_tensor(first_ptr)
-            rec = Store(tensor.data_ptr(), ptr.data - tensor.data_ptr(), mask.data)
-            rec.call_path = _extract_user_frames()
-            self.records.append(rec)
-
         def pre_masked_store_callback(ptr, keys, value, mask=None, *args, **kwargs):
             if not self.sample:
                 return

From 0c92fbdca40f75de8dd29c60419b40c4574b47c4 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Mon, 3 Nov 2025 14:30:55 -0500
Subject: [PATCH 68/90] more examples

---
 examples/nki/rmsnorm.py | 122 ++++++++++++++++++++++++++++++++++++++++
 examples/nki/softmax.py |  73 ++++++++++++++++++++++++
 triton_viz/core/nki.py  |  45 +++++++++++++++
 3 files changed, 240 insertions(+)
 create mode 100644 examples/nki/rmsnorm.py
 create mode 100644 examples/nki/softmax.py

diff --git a/examples/nki/rmsnorm.py b/examples/nki/rmsnorm.py
new file mode 100644
index 00000000..da87ec0c
--- /dev/null
+++ b/examples/nki/rmsnorm.py
@@ -0,0 +1,122 @@
+import neuronxcc.nki as nki
+import neuronxcc.nki.language as nl
+from triton_viz.clients import Tracer
+from triton_viz.core.trace import launches
+import math
+import numpy as np
+import torch
+import triton_viz
+
+
+def nki_rmsnorm_kernel(a_tensor, g_tensor, result):
+    # Calculate out_tensor = a_tensor/RMS(a_tensor) * g_tensor
+    # Where RMS(a_tensor) = sqrt((1/N) * sum(a_tensor * a_tensor))
+    # and N = a_tensor.shape[1]
+    # Reduction (mean) is performed in the free (2nd) dimension
+    B, D = a_tensor.shape
+
+    # Make sure shapes match
+    assert D == g_tensor.shape[0]
+
+    # Generate tensor indices to index input tensor
+    ix = nl.arange(2)[:, None]
+    iw = nl.arange(1)[:, None]
+    iy = nl.arange(8)[None, :]
+
+    # Load RMSNorm weight once, reused by rows/tiles of a_tensor
+    g_tile = nl.load(g_tensor.reshape((1, D))[iw, iy], mask=iy < D)
+
+    # Process 2 rows at a time due to 2-partition tile size limitation
+    # Since we're not reducing across the first dimension
+    # Tiles can be processed independently
+    for i in nl.affine_range(math.ceil(B / 2)):
+        # Load input data from external memory to on-chip memory
+        mask = (i * 2 + ix < B) & (iy < D)
+        a_tile = nl.load(a_tensor[i * 2 + ix, iy], mask=mask)
+
+        # Compute element-wise square of a_tensor
+        in_square = nl.square(a_tile)
+
+        # Calculate sum of squared elements, along last dimension
+        # square_sum = nl.sum(in_square, axis=[1], mask=mask) #[:, None]
+        square_sum = nl.sum(in_square, axis=1, mask=mask)[:, None]
+
+        # Scale and get a reciprocal
+        mean = square_sum / D
+
+        # Take square root of mean and then reciprocal with
+        # rsqrt API (one ISA instruction)
+        rms_reciprocal = nl.rsqrt(mean)
+
+        # Scale the input tensor
+        out_tile = nl.multiply(a_tile, rms_reciprocal)
+
+        # Broadcast weight along first axis to match tensor shape
+        # B_active = min(B - i * 2, 2)
+        g_bcast = g_tile.broadcast_to((2, 8))
+
+        # Multiply with the RMSNorm weight
+        out_tile = nl.multiply(out_tile, g_bcast, mask=(i * 2 + ix < B))
+
+        # store the addition results back to external memory (out_tensor)
+        nl.store(result[i * 2 + ix, iy], value=out_tile, mask=mask)
+
+
+# ref
+def torch_rmsnorm_kernel(a_tensor, g_tensor):
+    # Square the tensor (element-wise)
+    in_square = a_tensor.pow(2)
+    # Calculate means in the free dimension
+    mean = in_square.mean(dim=1, keepdim=True)
+    # Scale by reciprocal of sqrt(mean)
+    tensor = a_tensor * torch.rsqrt(mean)
+
+    # Scale the output by the weight
+    return tensor * g_tensor
+
+
+TRITON_VIZ = True
+kernel_grid = (1, 1, 1)
+a_tensor = torch.arange(15).float().view(3, 5)
+g_tensor = torch.arange(5).float()
+result = torch.empty_like(a_tensor).numpy()
+kernel_args = (a_tensor.numpy(), g_tensor.numpy(), result)
+
+if TRITON_VIZ:
+    print("Executing kernel with NKI interpreter...")
+    traced_kernel = triton_viz.trace(clients=Tracer(), backend="nki")(
+        nki_rmsnorm_kernel
+    )
+    kernel_instance = traced_kernel[kernel_grid]
+    kernel_instance(*kernel_args)
+
+    print(f"Number of launches: {len(launches)}")
+    if launches:
+        launch = launches[-1]
+        print(f"Number of records: {len(launch.records)}")
+        for i, record in enumerate(launch.records):
+            print(f"Record {i}: {type(record).__name__}")
+            if hasattr(record, "ptr"):
+                print(f"  ptr: {record.ptr}")
+            if hasattr(record, "offsets"):
+                print(f"  offsets shape: {record.offsets.shape}")
+            if hasattr(record, "masks"):
+                print(f"  masks shape: {record.masks.shape}")
+
+    # Try to launch visualization
+    try:
+        triton_viz.launch(share=False)
+    except Exception as e:
+        print(f"\nError during visualization: {e}")
+        import traceback
+
+        traceback.print_exc()
+else:
+    print("Executing NKI JIT-ed matmul_kernel...")
+    compiled_kernel = nki.jit(nki_rmsnorm_kernel, kernel_return=False)
+    z2 = nki.simulate_kernel(compiled_kernel[kernel_grid], *kernel_args)
+
+z2 = result
+z1 = torch_rmsnorm_kernel(a_tensor, g_tensor).numpy()
+print(np.max(np.abs(z1 - z2)))
+assert np.allclose(z1, z2)
diff --git a/examples/nki/softmax.py b/examples/nki/softmax.py
new file mode 100644
index 00000000..bd44b5e9
--- /dev/null
+++ b/examples/nki/softmax.py
@@ -0,0 +1,73 @@
+from neuronxcc import nki
+import neuronxcc.nki.language as nl
+
+import triton_viz
+from triton_viz.clients import Tracer
+from triton_viz.core.trace import launches
+import numpy as np
+import math
+
+
+def softmax_kernel(in_tensor, out_tensor):
+    """NKI softmax_kernel to compute softmax on the last dimension
+
+    Args:
+        in_tensor: an input tensor of shape [B,D], where B is a multiple of 128
+        out_tensor: the resulting output tensor of shape [B,D]
+    """
+    B, D = in_tensor.shape
+
+    # assert nl.tile_size.pmax == 128
+    num_tiles = math.ceil(B / 128)
+    for tile_idx in nl.affine_range(num_tiles):
+        i_p = tile_idx * 128 + nl.arange(128)[:, None]
+        i_f = nl.arange(D)[None, :]
+        mask = (i_p < B) | (i_f < 0)
+        tile = nl.load(in_tensor[i_p, i_f], mask=mask)
+        tile_exp = nl.exp(tile)
+        tile_expsum = nl.sum(tile_exp, -1, keepdims=True, mask=mask)
+        tile_softmax = tile_exp / tile_expsum
+        nl.store(out_tensor[i_p, i_f], value=tile_softmax, mask=mask)
+
+
+TRITON_VIZ = True
+kernel_grid = (1, 1, 1)
+x_small = np.random.rand(16, 32).astype(np.float32)
+y_small = np.empty(x_small.shape, dtype=x_small.dtype)
+kernel_args = (x_small, y_small)
+
+if TRITON_VIZ:
+    print("Executing softmax_kernel with NKI interpreter...")
+    traced_kernel = triton_viz.trace(clients=Tracer(), backend="nki")(softmax_kernel)
+    kernel_instance = traced_kernel[kernel_grid]
+    kernel_instance(*kernel_args)
+
+    print(f"Number of launches: {len(launches)}")
+    if launches:
+        launch = launches[-1]
+        print(f"Number of records: {len(launch.records)}")
+        for i, record in enumerate(launch.records):
+            print(f"Record {i}: {type(record).__name__}")
+            if hasattr(record, "ptr"):
+                print(f"  ptr: {record.ptr}")
+            if hasattr(record, "offsets"):
+                print(f"  offsets shape: {record.offsets.shape}")
+            if hasattr(record, "masks"):
+                print(f"  masks shape: {record.masks.shape}")
+
+    # Try to launch visualization
+    try:
+        triton_viz.launch(share=False)
+    except Exception as e:
+        print(f"\nError during visualization: {e}")
+        import traceback
+
+        traceback.print_exc()
+else:
+    print("Executing NKI JIT-ed softmax_kernel...")
+    compiled_kernel = nki.jit(softmax_kernel, kernel_return=False)
+    nki.simulate_kernel(compiled_kernel[kernel_grid], *kernel_args)
+
+y_expected = np.exp(x_small) / np.exp(x_small).sum(-1, keepdim=True)
+print(np.max(np.abs(y_expected - y_small)))
+assert np.allclose(y_expected, y_small)
diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index 1a46a340..c0860cd9 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -172,6 +172,18 @@ def __and__(self, other):
     def __or__(self, other):
         return self._binary_op(other, lambda a, b: a | b, "or", "|")
 
+    def reshape(self, *args, **kwargs):
+        return NDArray(
+            value=self.data.reshape(*args), name=f"{self.name}_reshape", **kwargs
+        )
+
+    def broadcast_to(self, *args, **kwargs):
+        return NDArray(
+            value=np.broadcast_to(self.data, *args),
+            name=f"{self.name}_broadcast_to",
+            **kwargs,
+        )
+
 
 class Builder:
     def __init__(self, grid_dims=None):
@@ -370,6 +382,35 @@ 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
+        return NDArray(
+            value=x.data.sum(*args, **kwargs), name=f"{x.name}_sum", **kwargs
+        )
+
+    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 range(self, stop):
         return range(stop)
 
@@ -399,6 +440,10 @@ def nki_patch_lang():
     nl.mgrid = NDArray(value=np.mgrid, buffer=nl.sbuf, name="mgrid")
     nl.matmul = nki_builder.matmul
     nl.copy = nki_builder.copy
+    nl.sum = nki_builder.sum
+    nl.square = nki_builder.square
+    nl.rsqrt = nki_builder.rsqrt
+    nl.multiply = nki_builder.multiply
 
     # attention-specific
     nl.load_transpose2d = nki_builder.load_transpose2d

From 0fc84eca9def99686eb32586c43a945858eab004 Mon Sep 17 00:00:00 2001
From: gujialiang123 <562960306@qq.com>
Date: Wed, 5 Nov 2025 05:15:16 -0500
Subject: [PATCH 69/90] apply ruff-format; move normalization to draw; stop
 tracking scripts/

---
 scripts/print_tracebacks.py        | 58 -----------------------
 triton_viz/visualizer/draw.py      | 74 ++++++++++++++++++++----------
 triton_viz/visualizer/interface.py | 74 +++++++++++++-----------------
 3 files changed, 80 insertions(+), 126 deletions(-)
 delete mode 100644 scripts/print_tracebacks.py

diff --git a/scripts/print_tracebacks.py b/scripts/print_tracebacks.py
deleted file mode 100644
index f6686eca..00000000
--- a/scripts/print_tracebacks.py
+++ /dev/null
@@ -1,58 +0,0 @@
-import os
-import sys
-import torch
-import triton
-from triton_viz.core import config as cfg
-from triton_viz.core.trace import launches
-from triton_viz.core.data import Load, Store, Dot
-
-# Ensure project root is on sys.path so we can import examples
-PROJECT_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), ".."))
-if PROJECT_ROOT not in sys.path:
-    sys.path.insert(0, PROJECT_ROOT)
-
-from examples.load_store import simple_kernel
-
-
-def main():
-    # Reset viz config to a clean state
-    cfg.reset()
-
-    # Prepare sample data and run the kernel once
-    device = "cpu"
-    size = 16
-    BLOCK_SIZE = 8
-    torch.manual_seed(0)
-    x = torch.arange(size, dtype=torch.float32, device=device)
-    output = torch.empty_like(x)
-    grid = lambda meta: (triton.cdiv(size, meta["BLOCK_SIZE"]),)
-
-    # Run the traced kernel (decorated in examples/load_store.py)
-    simple_kernel[grid](x, output, size, BLOCK_SIZE)
-
-    # Inspect the last launch and print a sample op's call_path
-    print(f"Number of launches: {len(launches)}")
-    if not launches:
-        return
-
-    launch = launches[-1]
-    print(f"Number of records: {len(launch.records)}")
-
-    target = None
-    for rec in launch.records:
-        if isinstance(rec, (Load, Store, Dot)):
-            target = rec
-            break
-
-    if target is None:
-        print("No Load/Store/Dot record found to display call_path.")
-        return
-
-    print(f"=== {type(target).__name__} call_path ===")
-    for i, f in enumerate(target.call_path):
-        line = f.line if hasattr(f, "line") else ""
-        print(f"{i:02d}  {f.filename}:{f.lineno}  [{f.name}]  |  {line}")
-
-
-if __name__ == "__main__":
-    main()
diff --git a/triton_viz/visualizer/draw.py b/triton_viz/visualizer/draw.py
index 90289230..be3bbd5b 100644
--- a/triton_viz/visualizer/draw.py
+++ b/triton_viz/visualizer/draw.py
@@ -182,9 +182,35 @@ def prepare_visualization_data(program_records, tensor_table):
                 }
             )
 
+            # Normalize Dot operands to NumPy arrays for downstream endpoints
+            try:
+                import numpy as _np
+
+                def _to_numpy_cpu(x):
+                    if isinstance(x, _np.ndarray):
+                        return x
+                    if hasattr(x, "cpu"):
+                        try:
+                            return x.detach().cpu().numpy()
+                        except Exception:
+                            return _np.asarray(x)
+                    if hasattr(x, "data"):
+                        return _np.asarray(getattr(x, "data"))
+                    if hasattr(x, "_value"):
+                        return _np.asarray(getattr(x, "_value"))
+                    return _np.asarray(x)
+
+                a_np = _to_numpy_cpu(record.input_data)
+                b_np = _to_numpy_cpu(record.other_data)
+            except Exception:
+                import numpy as _np
+
+                a_np = _np.asarray(record.input_data)
+                b_np = _np.asarray(record.other_data)
+
             raw_tensor_data[record_uuid] = {
-                "input_data": torch.tensor(record.input_data),
-                "other_data": torch.tensor(record.other_data),
+                "input_data": a_np,
+                "other_data": b_np,
                 "intermediate_results": record.intermediate_results,
                 "tracebacks": [
                     {
@@ -253,38 +279,36 @@ def prepare_visualization_data(program_records, tensor_table):
                 }
             )
 
-            # Normalize to torch.Tensor for downstream APIs
+            # Normalize to NumPy array for downstream APIs, and cache basic stats
             try:
                 import numpy as _np
-                import torch as _torch
 
                 gt = global_tensor.data
-                # gt could be torch.Tensor, numpy.ndarray, or NDArray wrapper
-                if (
-                    hasattr(gt, "cpu")
-                    and callable(getattr(gt, "cpu"))
-                    and hasattr(gt, "shape")
-                ):
-                    t_cpu = gt.cpu()
+                if hasattr(gt, "cpu") and callable(getattr(gt, "cpu", None)):
+                    try:
+                        arr = gt.detach().cpu().numpy()
+                    except Exception:
+                        arr = _np.asarray(gt)
+                elif hasattr(gt, "data"):
+                    arr = _np.asarray(getattr(gt, "data"))
+                elif hasattr(gt, "_value"):
+                    arr = _np.asarray(getattr(gt, "_value"))
                 else:
-                    # NDArray: prefer .data property -> numpy array
-                    if hasattr(gt, "data"):
-                        arr = gt.data
-                    else:
-                        arr = gt
-                    arr = _np.asarray(arr)
-                    t_cpu = _torch.from_numpy(
-                        arr.copy() if not arr.flags.c_contiguous else arr
-                    )
+                    arr = _np.asarray(gt)
             except Exception:
-                # Fallback to empty tensor if conversion fails
-                import torch as _torch
+                import numpy as _np
+
+                arr = _np.asarray([])
 
-                t_cpu = _torch.tensor([])
+            t_min = float(np.min(arr)) if arr.size else 0.0
+            t_max = float(np.max(arr)) if arr.size else 0.0
 
             raw_tensor_data[record_uuid] = {
-                "global_tensor": t_cpu,
-                "dims": len(t_cpu.shape),
+                "global_tensor": arr,
+                "dims": int(arr.ndim),
+                "shape": list(arr.shape),
+                "min": t_min,
+                "max": t_max,
                 "tracebacks": [
                     {
                         "filename": f.filename,
diff --git a/triton_viz/visualizer/interface.py b/triton_viz/visualizer/interface.py
index 1e300dae..f52cdb42 100644
--- a/triton_viz/visualizer/interface.py
+++ b/triton_viz/visualizer/interface.py
@@ -281,17 +281,19 @@ def get_matmul_c():
     if a is None or b is None:
         return jsonify({"error": "MatMul tensors not available"}), 200
     try:
-        # Compute C on CPU
-        c = a @ b
-        c_cpu = c.cpu()
-        cmin = float(c_cpu.min().item())
-        cmax = float(c_cpu.max().item())
+        import numpy as _np
+
+        a_np = _np.asarray(a)
+        b_np = _np.asarray(b)
+        c_np = a_np @ b_np
+        cmin = float(_np.min(c_np)) if c_np.size else 0.0
+        cmax = float(_np.max(c_np)) if c_np.size else 0.0
         return jsonify(
             {
-                "shape": list(c_cpu.shape),
+                "shape": list(c_np.shape),
                 "min": cmin,
                 "max": cmax,
-                "values": c_cpu.numpy().tolist(),
+                "values": c_np.tolist(),
             }
         )
     except Exception as e:
@@ -324,8 +326,12 @@ def get_matmul_vectors():
     if a is None or b is None:
         return jsonify({"error": "MatMul tensors not available"}), 200
     try:
-        a_row = a[row, :].cpu().numpy().tolist()
-        b_col = b[:, col].cpu().numpy().tolist()
+        import numpy as _np
+
+        a_np = _np.asarray(a)
+        b_np = _np.asarray(b)
+        a_row = a_np[row, :].tolist()
+        b_col = b_np[:, col].tolist()
         k = len(a_row)
         return jsonify({"row": row, "col": col, "a_row": a_row, "b_col": b_col, "k": k})
     except Exception as e:
@@ -396,29 +402,16 @@ def get_load_value():
         try:
             import numpy as _np
 
-            t = op_data["global_tensor"]
-            # Normalize to numpy array for robust indexing
-            if hasattr(t, "cpu"):
-                try:
-                    arr = t.cpu().numpy()
-                except Exception:
-                    arr = _np.asarray(t)
-            elif hasattr(t, "_value"):
-                arr = _np.asarray(getattr(t, "_value"))
-            elif hasattr(t, "data"):
-                arr = _np.asarray(getattr(t, "data"))
-            else:
-                arr = _np.asarray(t)
-
+            arr = _np.asarray(op_data["global_tensor"])  # already NumPy in draw.py
             yy, xx, zz = int(y), int(x), int(z)
-            value = 0.0
             if arr.ndim >= 3:
                 value = float(arr[yy, xx, zz])
             elif arr.ndim == 2:
                 value = float(arr[yy, xx])
             elif arr.ndim == 1:
                 value = float(arr[xx])
-
+            else:
+                value = 0.0
             return jsonify({"value": value})
         except IndexError:
             return jsonify({"error": "Coordinates out of bounds"}), 200
@@ -490,26 +483,21 @@ def get_load_tensor():
     import numpy as _np
 
     try:
-        t = op_data["global_tensor"]
-        if hasattr(t, "cpu"):
-            try:
-                arr = t.cpu().numpy()
-            except Exception:
-                arr = _np.asarray(t)
-        elif hasattr(t, "_value"):
-            arr = _np.asarray(getattr(t, "_value"))
-        elif hasattr(t, "data"):
-            arr = _np.asarray(getattr(t, "data"))
-        else:
-            arr = _np.asarray(t)
-        t_min = float(_np.min(arr)) if arr.size else 0.0
-        t_max = float(_np.max(arr)) if arr.size else 0.0
+        arr = _np.asarray(op_data["global_tensor"])  # already NumPy in draw.py
+        t_shape = op_data.get("shape") or list(arr.shape)
+        t_dims = op_data.get("dims") or int(arr.ndim)
+        t_min = op_data.get("min")
+        t_max = op_data.get("max")
+        if t_min is None:
+            t_min = float(_np.min(arr)) if arr.size else 0.0
+        if t_max is None:
+            t_max = float(_np.max(arr)) if arr.size else 0.0
         return jsonify(
             {
-                "shape": list(arr.shape),
-                "dims": int(arr.ndim),
-                "min": t_min,
-                "max": t_max,
+                "shape": t_shape,
+                "dims": int(t_dims),
+                "min": float(t_min),
+                "max": float(t_max),
                 "values": arr.tolist(),
             }
         )

From aceeae2757c3e47f1f0554f71c1bf5b0aa35961e Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Wed, 5 Nov 2025 13:59:26 -0500
Subject: [PATCH 70/90] Triton/NKI Load/Store unify pt 1. (push to same
 callback - not working)

---
 examples/load_store.py              |  2 +-
 triton_viz/clients/tracer/tracer.py | 62 ++++++++++++++++++-----------
 triton_viz/core/data.py             | 18 ++++++---
 triton_viz/core/patch.py            | 47 +++++++++++++++++-----
 4 files changed, 90 insertions(+), 39 deletions(-)

diff --git a/examples/load_store.py b/examples/load_store.py
index 24bd55e2..b3289f4d 100644
--- a/examples/load_store.py
+++ b/examples/load_store.py
@@ -50,7 +50,7 @@ def simple_kernel(
 
     # Try to launch visualization
     try:
-        triton_viz.launch()
+        triton_viz.launch(share=False)
     except Exception as e:
         print(f"\nError during visualization: {e}")
         import traceback
diff --git a/triton_viz/clients/tracer/tracer.py b/triton_viz/clients/tracer/tracer.py
index 5df8b588..b29b6b2d 100644
--- a/triton_viz/clients/tracer/tracer.py
+++ b/triton_viz/clients/tracer/tracer.py
@@ -7,8 +7,8 @@
     ReduceSum,
     Dot,
     Grid,
-    MaskedLoad,
-    MaskedStore,
+    # MaskedLoad,
+    # MaskedStore,
     RawLoad,
     RawStore,
     Allocate,
@@ -146,31 +146,47 @@ def _convert_keys_to_numpy(keys):
             else:
                 return keys
 
-        def pre_masked_load_callback(ptr, keys, mask=None, *args, **kwargs):
+        def pre_masked_load_callback(ptr, *args, **kwargs):
             if not self.sample:
                 return
-            keys = _convert_keys_to_numpy(keys)
 
-            self.records.append(
-                Load(
-                    ptr.data_ptr(),
-                    masked_load(ptr.get_offsets().data, keys, mask=mask.data),
-                    mask.data,
-                )
-            )
+            mask = kwargs["mask"]
+            if "keys" in kwargs:
+                keys = _convert_keys_to_numpy(kwargs["keys"])
+                offsets = masked_load(ptr.get_offsets().data, keys, mask=mask.data)
+                tensor = ptr
+            else:  # i.e. for triton, ptr = pointer + offsets
+                first_ptr = np.reshape(ptr.data, (-1))[0]
+                tensor = self._get_tensor(first_ptr)
+                offsets = ptr.data - tensor.data_ptr()
+
+            rec = Load(tensor.data_ptr(), offsets, mask.data)
+            rec.call_path = _extract_user_frames()
+            self.records.append(rec)
 
-        def pre_masked_store_callback(ptr, keys, value, mask=None, *args, **kwargs):
+        def pre_masked_store_callback(ptr, *args, **kwargs):
             if not self.sample:
                 return
-            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 = kwargs.get("mask", None)
+            if "keys" in kwargs:
+                keys = _convert_keys_to_numpy(kwargs["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)
+                tensor = ptr
+            else:  # i.e. for triton, ptr = pointer + offsets, so keys=None
+                first_ptr = np.reshape(ptr.data, (-1))[0]
+                tensor = self._get_tensor(first_ptr)
+                offsets = ptr.data - tensor.data_ptr()
                 mask_data = mask.data
-                offsets = masked_load(ptr.get_offsets().data, keys, mask=mask_data)
 
-            self.records.append(Store(ptr.data_ptr(), offsets, mask_data))
+            rec = Store(tensor.data_ptr(), offsets, mask_data)
+            rec.call_path = _extract_user_frames()
+            self.records.append(rec)
 
         # Raw (unmasked) ops: synthesize a full True mask based on ptr shape
         def pre_raw_load_callback(ptr):
@@ -238,13 +254,13 @@ def post_flip_callback(ret, x, *args, **kwargs):
 
         if op_type is Allocate:
             return OpCallbacks(after_callback=post_allocate_callback)
+        # elif op_type is Load:
+        #    return OpCallbacks(before_callback=pre_load_callback)
         elif op_type is Load:
-            return OpCallbacks(before_callback=pre_load_callback)
-        elif op_type is MaskedLoad:
             return OpCallbacks(before_callback=pre_masked_load_callback)
+        # elif op_type is Store:
+        #    return OpCallbacks(before_callback=pre_store_callback)
         elif op_type is Store:
-            return OpCallbacks(before_callback=pre_store_callback)
-        elif op_type is MaskedStore:
             return OpCallbacks(before_callback=pre_masked_store_callback)
         elif op_type is RawLoad:
             return OpCallbacks(before_callback=pre_raw_load_callback)
diff --git a/triton_viz/core/data.py b/triton_viz/core/data.py
index f278501b..7c73b54b 100644
--- a/triton_viz/core/data.py
+++ b/triton_viz/core/data.py
@@ -67,9 +67,12 @@ class Store(Op):
     time_idx: int = 0
 
 
-@dataclass
-class MaskedStore(Op):
-    name: ClassVar[str] = "masked_store"
+class TritonStore(Store):
+    ...
+
+
+class NKIStore(Store):
+    ...
 
 
 @dataclass
@@ -91,9 +94,12 @@ class Load(Op):
     time_idx: int = 0
 
 
-@dataclass
-class MaskedLoad(Op):
-    name: ClassVar[str] = "masked_load"
+class TritonLoad(Load):
+    ...
+
+
+class NKILoad(Load):
+    ...
 
 
 @dataclass
diff --git a/triton_viz/core/patch.py b/triton_viz/core/patch.py
index b97d00f9..293e1793 100644
--- a/triton_viz/core/patch.py
+++ b/triton_viz/core/patch.py
@@ -12,10 +12,8 @@
     Allocate,
     RawLoad,
     Load,
-    MaskedLoad,
-    RawStore,
     Store,
-    MaskedStore,
+    RawStore,
     UnaryOp,
     BinaryOp,
     TernaryOp,
@@ -60,12 +58,18 @@
 from triton.runtime import JITFunction
 from triton_viz.core.nki import nki_builder
 
+
 # shared operation registry for both backends
+def identity(*args, **kwargs):
+    return args, kwargs
+
+
 OPERATION_REGISTRY = {
     "triton": {
         "op_list": [
             ProgramId,
             RawStore,
+            # TritonStore,
             Store,
             RawLoad,
             Load,
@@ -102,10 +106,29 @@
         "original_ops": {
             ProgramId: interpreter_builder.create_get_program_id,
             RawStore: interpreter_builder.create_store,
+            # TritonStore: interpreter_builder.create_masked_store,
             Store: interpreter_builder.create_masked_store,
             RawLoad: interpreter_builder.create_load,
             Load: interpreter_builder.create_masked_load,
             Dot: interpreter_builder.create_dot,
+            # ProgramId: (
+            #    interpreter_builder.create_get_program_id,
+            #    lambda axis: axis
+            # ),
+            # RawStore: (interpreter_builder.create_store, identity),
+            # TritonStore: (
+            #    interpreter_builder.create_masked_store,
+            #    lambda ptr, value, mask, *args, **kwargs: ptr, value, mask
+            # ),
+            # RawLoad: interpreter_builder.create_load,
+            # TritonLoad: (
+            #    interpreter_builder.create_masked_load,
+            #    lambda ptr, mask, *args, **kwargs: ptr, mask
+            # ),
+            # Dot: (
+            #    interpreter_builder.create_dot,
+            #    lambda a, b, *args, **kwargs: a, b
+            # ),
             UnaryOp: interpreter_builder.unary_op,
             BinaryOp: interpreter_builder.binary_op,
             TernaryOp: interpreter_builder.ternary_op,
@@ -168,8 +191,10 @@
         "op_list": [
             Allocate,
             ProgramId,
-            MaskedStore,
-            MaskedLoad,
+            Load,
+            Store,
+            # NKILoad,
+            # NKIStore,
             Dot,
             UnaryOp,
             MakeRange,
@@ -177,8 +202,10 @@
         "original_ops": {
             ProgramId: nki_builder.program_id,
             Allocate: nki_builder.ndarray,
-            MaskedLoad: nki_builder.masked_load,
-            MaskedStore: nki_builder.masked_store,
+            # NKILoad: nki_builder.masked_load,
+            # NKIStore: nki_builder.masked_store,
+            Load: nki_builder.masked_load,
+            Store: nki_builder.masked_store,
             Dot: nki_builder.matmul,
             UnaryOp: nki_builder._unary_op,
             MakeRange: nki_builder.arange,
@@ -186,8 +213,10 @@
         "op_attr_names": {
             ProgramId: "program_id",
             Allocate: "ndarray",
-            MaskedLoad: "masked_load",
-            MaskedStore: "masked_store",
+            # NKILoad: "masked_load",
+            # NKIStore: "masked_store",
+            Load: "masked_load",
+            Store: "masked_store",
             Dot: "matmul",
             UnaryOp: "_unary_op",
             MakeRange: "arange",

From 10271aec8e7ceecd6a028ffae835e560255260f9 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Wed, 5 Nov 2025 23:13:02 -0500
Subject: [PATCH 71/90] Triton/NKI Load/Store unify pt. 2 (add adapters to
 standardize callback args for all clients)

---
 triton_viz/clients/profiler/profiler.py   |   6 +-
 triton_viz/clients/sanitizer/sanitizer.py |  10 +-
 triton_viz/clients/tracer/tracer.py       |  62 +--
 triton_viz/core/client.py                 |   4 +-
 triton_viz/core/patch.py                  | 437 +++++++++++++---------
 5 files changed, 292 insertions(+), 227 deletions(-)

diff --git a/triton_viz/clients/profiler/profiler.py b/triton_viz/clients/profiler/profiler.py
index f2b19bfa..eea48a9e 100644
--- a/triton_viz/clients/profiler/profiler.py
+++ b/triton_viz/clients/profiler/profiler.py
@@ -125,16 +125,14 @@ def _get_mask_stats(mask: TensorHandle) -> tuple[int, int]:
             false_count = np.count_nonzero(np.logical_not(mask.data))
             return total_count, false_count
 
-        def pre_load_callback(
-            ptr, mask, other, cache_modifier, eviction_policy, is_volatile
-        ):
+        def pre_load_callback(ptr, mask, keys):
             self._report_load_store_bytes("load", ptr, mask)
             if not self.disable_load_mask_percentage_check:
                 total_count, false_count = _get_mask_stats(mask)
                 self.load_mask_total_count += total_count
                 self.load_mask_false_count += false_count
 
-        def pre_store_callback(ptr, value, mask, cache_modifier, eviction_policy):
+        def pre_store_callback(ptr, mask, keys):
             self._report_load_store_bytes("store", ptr, mask)
             if not self.disable_load_mask_percentage_check:
                 total_count, false_count = _get_mask_stats(mask)
diff --git a/triton_viz/clients/sanitizer/sanitizer.py b/triton_viz/clients/sanitizer/sanitizer.py
index bc45d5f2..bcda870c 100644
--- a/triton_viz/clients/sanitizer/sanitizer.py
+++ b/triton_viz/clients/sanitizer/sanitizer.py
@@ -548,9 +548,7 @@ def grid_callback(self, grid: tuple[int, ...]) -> None:
         self.tensors = sorted(self.tensors, key=lambda x: x.data_ptr())
 
     def register_op_callback(self, op_type: type[Op], *args, **kwargs) -> OpCallbacks:
-        def pre_load_callback(
-            ptr, mask, other, cache_modifier, eviction_policy, is_volatile
-        ):
+        def pre_load_callback(ptr, mask, _keys):
             first_loc = np.unravel_index(np.argmax(mask, axis=None), mask.data.shape)
             first_ptr = ptr.data[first_loc]
             tensor = _get_tensor(self.tensors, first_ptr)
@@ -558,14 +556,12 @@ def pre_load_callback(
             self._report(op_type, oob)
             ptr.data = tensor.data_ptr() + oob["corrected_offsets"]
 
-        def pre_store_callback(ptr, value, mask, cache_modifier, eviction_policy):
+        def pre_store_callback(ptr, mask, _keys):
             first_loc = np.unravel_index(np.argmax(mask, axis=None), mask.data.shape)
             first_ptr = ptr.data[first_loc]
             tensor = _get_tensor(self.tensors, first_ptr)
             oob = check_out_of_bounds_access(ptr.data, mask.data, tensor)
-            self._report(
-                op_type, check_out_of_bounds_access(ptr.data, mask.data, tensor)
-            )
+            self._report(op_type, oob)
             ptr.data = tensor.data_ptr() + oob["corrected_offsets"]
 
         if op_type is Load:
diff --git a/triton_viz/clients/tracer/tracer.py b/triton_viz/clients/tracer/tracer.py
index b29b6b2d..af801221 100644
--- a/triton_viz/clients/tracer/tracer.py
+++ b/triton_viz/clients/tracer/tracer.py
@@ -7,8 +7,6 @@
     ReduceSum,
     Dot,
     Grid,
-    # MaskedLoad,
-    # MaskedStore,
     RawLoad,
     RawStore,
     Allocate,
@@ -115,28 +113,10 @@ def _extract_user_frames() -> list[traceback.FrameSummary]:
                     return [f]
             return stack[-1:]
 
-        def post_allocate_callback(ret, *args, **kwargs):
+        def post_allocate_callback(ret):
             assert hasattr(ret, "data")
             self.tensors.append(ret)
 
-        def pre_load_callback(ptr, mask, *args, **kwargs):
-            if not self.sample:
-                return
-            first_ptr = np.reshape(ptr.data, (-1))[0]
-            tensor = self._get_tensor(first_ptr)
-            rec = Load(tensor.data_ptr(), ptr.data - tensor.data_ptr(), mask.data)
-            rec.call_path = _extract_user_frames()
-            self.records.append(rec)
-
-        def pre_store_callback(ptr, value, mask, *args, **kwargs):
-            if not self.sample:
-                return
-            first_ptr = np.reshape(ptr.data, (-1))[0]
-            tensor = self._get_tensor(first_ptr)
-            rec = Store(tensor.data_ptr(), ptr.data - tensor.data_ptr(), mask.data)
-            rec.call_path = _extract_user_frames()
-            self.records.append(rec)
-
         def _convert_keys_to_numpy(keys):
             """Convert any NDArrays in keys to numpy arrays."""
             if isinstance(keys, (tuple, list)):
@@ -146,31 +126,34 @@ def _convert_keys_to_numpy(keys):
             else:
                 return keys
 
-        def pre_masked_load_callback(ptr, *args, **kwargs):
+        def pre_masked_load_callback(ptr, mask, keys):
             if not self.sample:
                 return
 
-            mask = kwargs["mask"]
-            if "keys" in kwargs:
-                keys = _convert_keys_to_numpy(kwargs["keys"])
-                offsets = masked_load(ptr.get_offsets().data, keys, mask=mask.data)
-                tensor = ptr
-            else:  # i.e. for triton, ptr = pointer + offsets
+            if keys is None:  # i.e. for triton, ptr = pointer + offsets
                 first_ptr = np.reshape(ptr.data, (-1))[0]
                 tensor = self._get_tensor(first_ptr)
                 offsets = ptr.data - tensor.data_ptr()
+            else:
+                keys = _convert_keys_to_numpy(keys)
+                offsets = masked_load(ptr.get_offsets().data, keys, mask=mask.data)
+                tensor = ptr
 
             rec = Load(tensor.data_ptr(), offsets, mask.data)
             rec.call_path = _extract_user_frames()
             self.records.append(rec)
 
-        def pre_masked_store_callback(ptr, *args, **kwargs):
+        def pre_masked_store_callback(ptr, mask, keys):
             if not self.sample:
                 return
 
-            mask = kwargs.get("mask", None)
-            if "keys" in kwargs:
-                keys = _convert_keys_to_numpy(kwargs["keys"])
+            if keys is None:  # i.e. for triton, ptr = pointer + offsets, so keys=None
+                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)
                 if mask is None:
                     offsets = masked_load(ptr.get_offsets().data, keys)
                     mask_data = np.ones_like(offsets).astype(bool)
@@ -178,11 +161,6 @@ def pre_masked_store_callback(ptr, *args, **kwargs):
                     mask_data = mask.data
                     offsets = masked_load(ptr.get_offsets().data, keys, mask=mask_data)
                 tensor = ptr
-            else:  # i.e. for triton, ptr = pointer + offsets, so keys=None
-                first_ptr = np.reshape(ptr.data, (-1))[0]
-                tensor = self._get_tensor(first_ptr)
-                offsets = ptr.data - tensor.data_ptr()
-                mask_data = mask.data
 
             rec = Store(tensor.data_ptr(), offsets, mask_data)
             rec.call_path = _extract_user_frames()
@@ -211,16 +189,14 @@ def pre_raw_store_callback(ptr, value):
             rec.call_path = _extract_user_frames()
             self.records.append(rec)
 
-        def post_reduce_sum_callback(
-            ret, input, axis=None, keep_dims=False, *args, **kwargs
-        ):
+        def post_reduce_sum_callback(ret, input, axis, keep_dims):
             if not self.sample:
                 return
             input_shape = input.handle.data.shape
             output_shape = ret.handle.data.shape
             self.records.append(ReduceSum(input_shape, axis, keep_dims, output_shape))
 
-        def post_dot_callback(ret, input, other, *args, **kwargs):
+        def post_dot_callback(ret, input, other):
             if not self.sample:
                 return
             input_shape = input.data.shape
@@ -254,12 +230,8 @@ def post_flip_callback(ret, x, *args, **kwargs):
 
         if op_type is Allocate:
             return OpCallbacks(after_callback=post_allocate_callback)
-        # elif op_type is Load:
-        #    return OpCallbacks(before_callback=pre_load_callback)
         elif op_type is Load:
             return OpCallbacks(before_callback=pre_masked_load_callback)
-        # elif op_type is Store:
-        #    return OpCallbacks(before_callback=pre_store_callback)
         elif op_type is Store:
             return OpCallbacks(before_callback=pre_masked_store_callback)
         elif op_type is RawLoad:
diff --git a/triton_viz/core/client.py b/triton_viz/core/client.py
index d1254257..fbbbedad 100644
--- a/triton_viz/core/client.py
+++ b/triton_viz/core/client.py
@@ -130,11 +130,11 @@ def patch_run(self, fn, backend: str):
         with patch_calls(backend):
             for client in self.clients.values():
                 # get operations for the specified backend
-                backend_ops = OPERATION_REGISTRY[backend]["op_list"]
+                backend_ops: list = OPERATION_REGISTRY[backend]["op_list"]
 
                 for op in backend_ops:
                     # patch ops
-                    callbacks = client.register_op_callback(op, backend=backend)
+                    callbacks = client.register_op_callback(op)
                     patch_op(op, callbacks, backend=backend)
 
                 # patch for loops
diff --git a/triton_viz/core/patch.py b/triton_viz/core/patch.py
index 293e1793..47dfe936 100644
--- a/triton_viz/core/patch.py
+++ b/triton_viz/core/patch.py
@@ -1,6 +1,7 @@
 import triton.language as tl
 from contextlib import contextmanager
 from collections.abc import Callable
+from dataclasses import dataclass
 from typing import Any, Optional
 from functools import partialmethod
 from tqdm import tqdm
@@ -59,176 +60,266 @@
 from triton_viz.core.nki import nki_builder
 
 
-# shared operation registry for both backends
-def identity(*args, **kwargs):
-    return args, kwargs
+@dataclass
+class AdapterResult:
+    """
+    For each backend, ops may have slightly different function signatures
+    which we run through (backend, function)-specific adapters to return
+    standardized args/kwargs for client callbacks.
+    """
+
+    args: tuple[Any, ...]
+    kwargs: dict[str, Any]
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
+        self.args = args
+        self.kwargs = kwargs
+
+
+def passthrough_adapter(*args: Any, **kwargs: Any) -> AdapterResult:
+    """Return arguments unchanged for clients that expect the original signature."""
+    return AdapterResult(args, kwargs)
+
+
+def _program_id_adapter(axis: Any, *_args: Any, **_kwargs: Any) -> AdapterResult:
+    return AdapterResult(axis)
+
+
+def _triton_raw_store_adapter(
+    ptr: Any, value: Any, *_args: Any, **_kwargs: Any
+) -> AdapterResult:
+    return AdapterResult(ptr, value)
+
+
+def _triton_store_adapter(
+    ptr: Any, _value: Any, mask: Any, *_args: Any, **kwargs: Any
+) -> AdapterResult:
+    keys = kwargs.get("keys")
+    return AdapterResult(ptr, mask, keys)
+
+
+def _triton_raw_load_adapter(ptr: Any, *_args: Any, **_kwargs: Any) -> AdapterResult:
+    return AdapterResult(ptr)
+
+
+def _triton_load_adapter(
+    ptr: Any, mask: Any, _other: Any, *_args: Any, **kwargs: Any
+) -> AdapterResult:
+    keys = kwargs.get("keys")
+    return AdapterResult(ptr, mask, keys)
+
+
+def _triton_dot_adapter(a: Any, b: Any, *_args: Any, **_kwargs: Any) -> AdapterResult:
+    return AdapterResult(a, b)
+
+
+def _triton_reduce_sum_adapter(
+    input_tensor, axis=None, keep_dims=False, *_args, **_kwargs
+) -> AdapterResult:
+    return AdapterResult(input_tensor, axis, keep_dims)
+
+
+def _triton_addptr_adapter(
+    ptr: Any, offset: Any, *_args: Any, **_kwargs: Any
+) -> AdapterResult:
+    return AdapterResult(ptr, offset)
+
+
+def _nki_load_adapter(
+    src: Any, keys: Any, *, mask: Optional[Any] = None, **_kwargs: Any
+) -> AdapterResult:
+    return AdapterResult(src, mask, keys)
+
+
+def _nki_store_adapter(
+    dst: Any, keys: Any, value: Any, *, mask: Optional[Any] = None, **_kwargs: Any
+) -> AdapterResult:
+    return AdapterResult(dst, mask, keys)
+
+
+def _nki_dot_adapter(x: Any, y: Any, *_args: Any, **_kwargs: Any) -> AdapterResult:
+    return AdapterResult(x, y)
+
+
+TRITON_OP_LIST = [
+    ProgramId,
+    RawStore,
+    Store,
+    RawLoad,
+    Load,
+    UnaryOp,
+    BinaryOp,
+    TernaryOp,
+    Dot,
+    MakeRange,
+    AddPtr,
+    Splat,
+    ExpandDims,
+    Broadcast,
+    ReduceMax,
+    ReduceMin,
+    ReduceSum,
+    MakeBlockPointer,
+    TensorPointerLoad,
+    TensorPointerStore,
+    Idiv,
+    Rsqrt,
+    CastImpl,
+    Reshape,
+    Join,
+    Fabs,
+    Ashr,
+    Advance,
+    FpToFp,
+    Umulhi,
+    Trans,
+    CumSum,
+    Bitcast,
+    AtomicCas,
+]
+
+TRITON_ORIGINAL_OPS = {
+    ProgramId: interpreter_builder.create_get_program_id,
+    RawStore: interpreter_builder.create_store,
+    Store: interpreter_builder.create_masked_store,
+    RawLoad: interpreter_builder.create_load,
+    Load: interpreter_builder.create_masked_load,
+    Dot: interpreter_builder.create_dot,
+    UnaryOp: interpreter_builder.unary_op,
+    BinaryOp: interpreter_builder.binary_op,
+    TernaryOp: interpreter_builder.ternary_op,
+    MakeRange: interpreter_builder.create_make_range,
+    AddPtr: interpreter_builder.create_addptr,
+    ExpandDims: interpreter_builder.create_expand_dims,
+    Broadcast: interpreter_builder.create_broadcast,
+    Splat: interpreter_builder.create_splat,
+    MakeBlockPointer: interpreter_builder.create_make_block_ptr,
+    TensorPointerLoad: interpreter_builder.create_tensor_pointer_load,
+    TensorPointerStore: interpreter_builder.create_tensor_pointer_store,
+    Idiv: interpreter_builder.create_idiv,
+    Rsqrt: interpreter_builder.create_rsqrt,
+    CastImpl: interpreter_builder.cast_impl,
+    Reshape: interpreter_builder.create_reshape,
+    Join: interpreter_builder.create_join,
+    Fabs: interpreter_builder.create_fabs,
+    Ashr: interpreter_builder.create_ashr,
+    Advance: interpreter_builder.create_advance,
+    FpToFp: interpreter_builder.create_fp_to_fp,
+    Umulhi: interpreter_builder.create_umulhi,
+    Trans: interpreter_builder.create_trans,
+    Bitcast: interpreter_builder.create_bitcast,
+    AtomicCas: interpreter_builder.create_atomic_cas,
+}
+
+TRITON_OP_ATTR_NAMES = {
+    ProgramId: "create_get_program_id",
+    RawStore: "create_store",
+    Store: "create_masked_store",
+    RawLoad: "create_load",
+    Load: "create_masked_load",
+    Dot: "create_dot",
+    UnaryOp: "unary_op",
+    BinaryOp: "binary_op",
+    TernaryOp: "ternary_op",
+    MakeRange: "create_make_range",
+    AddPtr: "create_addptr",
+    ExpandDims: "create_expand_dims",
+    Broadcast: "create_broadcast",
+    Splat: "create_splat",
+    MakeBlockPointer: "create_make_block_ptr",
+    TensorPointerLoad: "create_tensor_pointer_load",
+    TensorPointerStore: "create_tensor_pointer_store",
+    Idiv: "create_idiv",
+    Rsqrt: "create_rsqrt",
+    CastImpl: "cast_impl",
+    Reshape: "create_reshape",
+    Join: "create_join",
+    Fabs: "create_fabs",
+    Ashr: "create_ashr",
+    Advance: "create_advance",
+    FpToFp: "create_fp_to_fp",
+    Umulhi: "create_umulhi",
+    Trans: "create_trans",
+    Bitcast: "create_bitcast",
+    AtomicCas: "create_atomic_cas",
+}
+
+TRITON_ADAPTERS: dict[type[Op], Callable[..., AdapterResult]] = {
+    ProgramId: _program_id_adapter,
+    RawStore: _triton_raw_store_adapter,
+    Store: _triton_store_adapter,
+    RawLoad: _triton_raw_load_adapter,
+    Load: _triton_load_adapter,
+    Dot: _triton_dot_adapter,
+    ReduceSum: _triton_reduce_sum_adapter,
+    AddPtr: _triton_addptr_adapter,
+}
+
+for op_type in TRITON_OP_LIST:
+    TRITON_ADAPTERS.setdefault(op_type, passthrough_adapter)
+
+
+NKI_OP_LIST = [
+    Allocate,
+    ProgramId,
+    Load,
+    Store,
+    Dot,
+    UnaryOp,
+    MakeRange,
+]
+
+NKI_ORIGINAL_OPS = {
+    ProgramId: nki_builder.program_id,
+    Allocate: nki_builder.ndarray,
+    Load: nki_builder.masked_load,
+    Store: nki_builder.masked_store,
+    Dot: nki_builder.matmul,
+    UnaryOp: nki_builder._unary_op,
+    MakeRange: nki_builder.arange,
+}
+
+NKI_OP_ATTR_NAMES = {
+    ProgramId: "program_id",
+    Allocate: "ndarray",
+    Load: "masked_load",
+    Store: "masked_store",
+    Dot: "matmul",
+    UnaryOp: "_unary_op",
+    MakeRange: "arange",
+}
+
+NKI_ADAPTERS: dict[type[Op], Callable[..., AdapterResult]] = {
+    ProgramId: _program_id_adapter,
+    Load: _nki_load_adapter,
+    Store: _nki_store_adapter,
+    Dot: _nki_dot_adapter,
+}
+
+for op_type in NKI_OP_LIST:
+    NKI_ADAPTERS.setdefault(op_type, passthrough_adapter)
 
 
-OPERATION_REGISTRY = {
+OPERATION_REGISTRY: dict[str, dict[str, Any]] = {
     "triton": {
-        "op_list": [
-            ProgramId,
-            RawStore,
-            # TritonStore,
-            Store,
-            RawLoad,
-            Load,
-            UnaryOp,
-            BinaryOp,
-            TernaryOp,
-            Dot,
-            MakeRange,
-            AddPtr,
-            Splat,
-            ExpandDims,
-            Broadcast,
-            ReduceMax,
-            ReduceMin,
-            ReduceSum,
-            MakeBlockPointer,
-            TensorPointerLoad,
-            TensorPointerStore,
-            Idiv,
-            Rsqrt,
-            CastImpl,
-            Reshape,
-            Join,
-            Fabs,
-            Ashr,
-            Advance,
-            FpToFp,
-            Umulhi,
-            Trans,
-            CumSum,
-            Bitcast,
-            AtomicCas,
-        ],
-        "original_ops": {
-            ProgramId: interpreter_builder.create_get_program_id,
-            RawStore: interpreter_builder.create_store,
-            # TritonStore: interpreter_builder.create_masked_store,
-            Store: interpreter_builder.create_masked_store,
-            RawLoad: interpreter_builder.create_load,
-            Load: interpreter_builder.create_masked_load,
-            Dot: interpreter_builder.create_dot,
-            # ProgramId: (
-            #    interpreter_builder.create_get_program_id,
-            #    lambda axis: axis
-            # ),
-            # RawStore: (interpreter_builder.create_store, identity),
-            # TritonStore: (
-            #    interpreter_builder.create_masked_store,
-            #    lambda ptr, value, mask, *args, **kwargs: ptr, value, mask
-            # ),
-            # RawLoad: interpreter_builder.create_load,
-            # TritonLoad: (
-            #    interpreter_builder.create_masked_load,
-            #    lambda ptr, mask, *args, **kwargs: ptr, mask
-            # ),
-            # Dot: (
-            #    interpreter_builder.create_dot,
-            #    lambda a, b, *args, **kwargs: a, b
-            # ),
-            UnaryOp: interpreter_builder.unary_op,
-            BinaryOp: interpreter_builder.binary_op,
-            TernaryOp: interpreter_builder.ternary_op,
-            MakeRange: interpreter_builder.create_make_range,
-            AddPtr: interpreter_builder.create_addptr,
-            ExpandDims: interpreter_builder.create_expand_dims,
-            Broadcast: interpreter_builder.create_broadcast,
-            Splat: interpreter_builder.create_splat,
-            MakeBlockPointer: interpreter_builder.create_make_block_ptr,
-            TensorPointerLoad: interpreter_builder.create_tensor_pointer_load,
-            TensorPointerStore: interpreter_builder.create_tensor_pointer_store,
-            Idiv: interpreter_builder.create_idiv,
-            Rsqrt: interpreter_builder.create_rsqrt,
-            CastImpl: interpreter_builder.cast_impl,
-            Reshape: interpreter_builder.create_reshape,
-            Join: interpreter_builder.create_join,
-            Fabs: interpreter_builder.create_fabs,
-            Ashr: interpreter_builder.create_ashr,
-            Advance: interpreter_builder.create_advance,
-            FpToFp: interpreter_builder.create_fp_to_fp,
-            Umulhi: interpreter_builder.create_umulhi,
-            Trans: interpreter_builder.create_trans,
-            Bitcast: interpreter_builder.create_bitcast,
-            AtomicCas: interpreter_builder.create_atomic_cas,
-        },
-        "op_attr_names": {
-            ProgramId: "create_get_program_id",
-            RawStore: "create_store",
-            Store: "create_masked_store",
-            RawLoad: "create_load",
-            Load: "create_masked_load",
-            Dot: "create_dot",
-            UnaryOp: "unary_op",
-            BinaryOp: "binary_op",
-            TernaryOp: "ternary_op",
-            MakeRange: "create_make_range",
-            AddPtr: "create_addptr",
-            ExpandDims: "create_expand_dims",
-            Broadcast: "create_broadcast",
-            Splat: "create_splat",
-            MakeBlockPointer: "create_make_block_ptr",
-            TensorPointerLoad: "create_tensor_pointer_load",
-            TensorPointerStore: "create_tensor_pointer_store",
-            Idiv: "create_idiv",
-            Rsqrt: "create_rsqrt",
-            CastImpl: "cast_impl",
-            Reshape: "create_reshape",
-            Join: "create_join",
-            Fabs: "create_fabs",
-            Ashr: "create_ashr",
-            Advance: "create_advance",
-            FpToFp: "create_fp_to_fp",
-            Umulhi: "create_umulhi",
-            Trans: "create_trans",
-            Bitcast: "create_bitcast",
-            AtomicCas: "create_atomic_cas",
-        },
+        "op_list": TRITON_OP_LIST,
+        "original_ops": TRITON_ORIGINAL_OPS,
+        "op_attr_names": TRITON_OP_ATTR_NAMES,
+        "adapters": TRITON_ADAPTERS,
     },
     "nki": {
-        "op_list": [
-            Allocate,
-            ProgramId,
-            Load,
-            Store,
-            # NKILoad,
-            # NKIStore,
-            Dot,
-            UnaryOp,
-            MakeRange,
-        ],
-        "original_ops": {
-            ProgramId: nki_builder.program_id,
-            Allocate: nki_builder.ndarray,
-            # NKILoad: nki_builder.masked_load,
-            # NKIStore: nki_builder.masked_store,
-            Load: nki_builder.masked_load,
-            Store: nki_builder.masked_store,
-            Dot: nki_builder.matmul,
-            UnaryOp: nki_builder._unary_op,
-            MakeRange: nki_builder.arange,
-        },
-        "op_attr_names": {
-            ProgramId: "program_id",
-            Allocate: "ndarray",
-            # NKILoad: "masked_load",
-            # NKIStore: "masked_store",
-            Load: "masked_load",
-            Store: "masked_store",
-            Dot: "matmul",
-            UnaryOp: "_unary_op",
-            MakeRange: "arange",
-        },
+        "op_list": NKI_OP_LIST,
+        "original_ops": NKI_ORIGINAL_OPS,
+        "op_attr_names": NKI_OP_ATTR_NAMES,
+        "adapters": NKI_ADAPTERS,
     },
 }
 
 BUILDER = interpreter_builder
 current_backend = "triton"
-op_list = OPERATION_REGISTRY[current_backend]["op_list"]
-original_ops = OPERATION_REGISTRY[current_backend]["original_ops"]  # type: ignore
-_OP_ATTR_NAMES = OPERATION_REGISTRY[current_backend]["op_attr_names"]  # type: ignore
+op_list: list = OPERATION_REGISTRY[current_backend]["op_list"]
+original_ops: dict = OPERATION_REGISTRY[current_backend]["original_ops"]
+_OP_ATTR_NAMES: dict = OPERATION_REGISTRY[current_backend]["op_attr_names"]
 
 
 def get_builder_for_backend(backend: str):
@@ -272,14 +363,17 @@ def __init__(
         op: Callable,
         op_type: type[Op],
         callbacks: OpCallbacks,
+        adapter: Callable[..., AdapterResult],
     ):
         self.op = op
         self.op_type = op_type
         self.callbacks = callbacks
+        self.adapter = adapter
 
     def __call__(self, *args, **kwargs):
         if self.callbacks.before_callback:
-            self.callbacks.before_callback(*args, **kwargs)
+            before_args = self.adapter(*args, **kwargs)
+            self.callbacks.before_callback(*before_args.args, **before_args.kwargs)
         if self.callbacks.op_overrider:
             if self.op_type in reduce_map or self.op_type in scan_map:
                 # see triton.runtime.interpreter:ReduceOps.sum
@@ -301,7 +395,8 @@ def __call__(self, *args, **kwargs):
             ret = self.op(*args, **kwargs)
         if self.callbacks.after_callback:
             # Pass ret so that we don't have to derive output shape from args
-            self.callbacks.after_callback(ret, *args, **kwargs)
+            after_args = self.adapter(*args, **kwargs)
+            self.callbacks.after_callback(ret, *after_args.args, **after_args.kwargs)
         return ret
 
 
@@ -319,14 +414,16 @@ def patch_op(op_type: type[Op], callbacks: OpCallbacks, backend: Optional[str] =
     if backend not in OPERATION_REGISTRY:
         raise ValueError(f"Unknown backend: {backend}")
 
-    backend_ops = OPERATION_REGISTRY[backend]["original_ops"]  # type: ignore
-    backend_attr_names = OPERATION_REGISTRY[backend]["op_attr_names"]  # type: ignore
+    backend_ops = OPERATION_REGISTRY[backend]["original_ops"]
+    backend_attr_names = OPERATION_REGISTRY[backend]["op_attr_names"]
+    backend_adapters = OPERATION_REGISTRY[backend]["adapters"]
     backend_builder = nki_builder if backend == "nki" else interpreter_builder
 
     if op_type in backend_ops:
-        op_name = backend_attr_names[op_type]  # type: ignore
-        original_op = backend_ops[op_type]  # type: ignore
-        patched_op = PatchOp(original_op, op_type, callbacks)
+        op_name = backend_attr_names[op_type]
+        original_op = backend_ops[op_type]
+        adapter = backend_adapters[op_type]
+        patched_op = PatchOp(original_op, op_type, callbacks, adapter)
         setattr(
             backend_builder,
             op_name,
@@ -338,12 +435,14 @@ def patch_op(op_type: type[Op], callbacks: OpCallbacks, backend: Optional[str] =
         elif op_type in scan_map:
             op_name = scan_map[op_type].__name__
         original_op = getattr(tl, op_name)
-        patched_op = PatchOp(original_op, op_type, callbacks)
+        adapter = backend_adapters[op_type]
+        patched_op = PatchOp(original_op, op_type, callbacks, adapter)
         setattr(tl, op_name, lambda *args, **kwargs: patched_op(*args, **kwargs))
     elif backend == "triton" and op_type in math_map:
         op_name = math_map[op_type].__name__
         original_op = getattr(tl.math, op_name)
-        patched_op = PatchOp(original_op, op_type, callbacks)
+        adapter = backend_adapters[op_type]
+        patched_op = PatchOp(original_op, op_type, callbacks, adapter)
         setattr(tl.math, op_name, lambda *args, **kwargs: patched_op(*args, **kwargs))
     else:
         raise ValueError(f"Patching operator {op_type} not supported")

From 9f37b245ec611a25996fd1df34ebe8dad0c0a593 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Wed, 5 Nov 2025 23:20:32 -0500
Subject: [PATCH 72/90] whoops forgot to add allocate adapter

---
 triton_viz/core/client.py |  2 +-
 triton_viz/core/patch.py  | 11 ++++++++---
 2 files changed, 9 insertions(+), 4 deletions(-)

diff --git a/triton_viz/core/client.py b/triton_viz/core/client.py
index fbbbedad..d4545980 100644
--- a/triton_viz/core/client.py
+++ b/triton_viz/core/client.py
@@ -130,7 +130,7 @@ def patch_run(self, fn, backend: str):
         with patch_calls(backend):
             for client in self.clients.values():
                 # get operations for the specified backend
-                backend_ops: list = OPERATION_REGISTRY[backend]["op_list"]
+                backend_ops: list[type[Op]] = OPERATION_REGISTRY[backend]["op_list"]
 
                 for op in backend_ops:
                     # patch ops
diff --git a/triton_viz/core/patch.py b/triton_viz/core/patch.py
index 47dfe936..c38a5b80 100644
--- a/triton_viz/core/patch.py
+++ b/triton_viz/core/patch.py
@@ -125,6 +125,10 @@ def _triton_addptr_adapter(
     return AdapterResult(ptr, offset)
 
 
+def _nki_allocate_adapter(*_args: Any, **_kwargs: Any) -> AdapterResult:
+    return AdapterResult()
+
+
 def _nki_load_adapter(
     src: Any, keys: Any, *, mask: Optional[Any] = None, **_kwargs: Any
 ) -> AdapterResult:
@@ -291,6 +295,7 @@ def _nki_dot_adapter(x: Any, y: Any, *_args: Any, **_kwargs: Any) -> AdapterResu
 
 NKI_ADAPTERS: dict[type[Op], Callable[..., AdapterResult]] = {
     ProgramId: _program_id_adapter,
+    Allocate: _nki_allocate_adapter,
     Load: _nki_load_adapter,
     Store: _nki_store_adapter,
     Dot: _nki_dot_adapter,
@@ -317,9 +322,9 @@ def _nki_dot_adapter(x: Any, y: Any, *_args: Any, **_kwargs: Any) -> AdapterResu
 
 BUILDER = interpreter_builder
 current_backend = "triton"
-op_list: list = OPERATION_REGISTRY[current_backend]["op_list"]
-original_ops: dict = OPERATION_REGISTRY[current_backend]["original_ops"]
-_OP_ATTR_NAMES: dict = OPERATION_REGISTRY[current_backend]["op_attr_names"]
+op_list: list[Op] = OPERATION_REGISTRY[current_backend]["op_list"]
+original_ops: dict[Op, Callable] = OPERATION_REGISTRY[current_backend]["original_ops"]
+_OP_ATTR_NAMES: dict[Op, str] = OPERATION_REGISTRY[current_backend]["op_attr_names"]
 
 
 def get_builder_for_backend(backend: str):

From cd623f5d92d4e3f2dd147f9c490dcf2f27ea59d1 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Wed, 5 Nov 2025 23:47:17 -0500
Subject: [PATCH 73/90] be gone

---
 triton_viz/core/data.py | 16 ----------------
 1 file changed, 16 deletions(-)

diff --git a/triton_viz/core/data.py b/triton_viz/core/data.py
index 7c73b54b..efe8ca50 100644
--- a/triton_viz/core/data.py
+++ b/triton_viz/core/data.py
@@ -67,14 +67,6 @@ class Store(Op):
     time_idx: int = 0
 
 
-class TritonStore(Store):
-    ...
-
-
-class NKIStore(Store):
-    ...
-
-
 @dataclass
 class RawLoad(Op):
     name: ClassVar[str] = "raw_load"
@@ -94,14 +86,6 @@ class Load(Op):
     time_idx: int = 0
 
 
-class TritonLoad(Load):
-    ...
-
-
-class NKILoad(Load):
-    ...
-
-
 @dataclass
 class UnaryOp(Op):
     name: ClassVar[str] = "unary_op"

From 6470e5b335a5193a046606c624c0f6f0888d173a Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Thu, 6 Nov 2025 00:02:02 -0500
Subject: [PATCH 74/90] add adapter tests

---
 tests/test_adapters.py   | 168 +++++++++++++++++++++++++++++++++++++++
 triton_viz/core/patch.py |   5 +-
 2 files changed, 170 insertions(+), 3 deletions(-)
 create mode 100644 tests/test_adapters.py

diff --git a/tests/test_adapters.py b/tests/test_adapters.py
new file mode 100644
index 00000000..f1345cb0
--- /dev/null
+++ b/tests/test_adapters.py
@@ -0,0 +1,168 @@
+from triton_viz.core.callbacks import OpCallbacks
+from triton_viz.core.data import (
+    AddPtr,
+    Load,
+    ProgramId,
+    ReduceSum,
+    Store,
+)
+from triton_viz.core.patch import (
+    AdapterResult,
+    NKI_ADAPTERS,
+    PatchOp,
+    TRITON_ADAPTERS,
+)
+
+
+def test_adapter_result_kwargs_copy():
+    """Test that AdapterResult SHALLOW-copies kwargs."""
+    kwargs = {"named": "value"}
+    result = AdapterResult(1, **kwargs)
+    kwargs[
+        "named"
+    ] = "updated"  # result.kwargs["named"] is immutable str so it doesn't change from this
+    assert result.args == (1,)
+    assert result.kwargs == {"named": "value"}
+
+    mutable_object = ["before"]
+    kwargs = {"named": mutable_object}
+    result = AdapterResult(1, **kwargs)
+    kwargs["named"][
+        0
+    ] = "after"  # result.kwargs["named"] is mutable so it DOES change from this
+    assert result.args == (1,)
+    assert result.kwargs["named"][0] == "after"
+
+
+def test_triton_store_adapter_handles_keys():
+    """Adapter mirrors tl.store(ptr, value, mask=..., cache_modifier=..., eviction_policy=..., keys=...)."""
+    adapter = TRITON_ADAPTERS[Store]
+    ptr, value, mask, keys = object(), object(), object(), object()
+
+    cache_modifier = "wb"
+    eviction_policy = "evict_last"
+
+    full = adapter(
+        ptr,
+        value,
+        mask,
+        cache_modifier=cache_modifier,
+        eviction_policy=eviction_policy,
+        keys=keys,
+    )
+    assert isinstance(full, AdapterResult)
+    assert full.args == (ptr, mask, keys)
+    assert full.kwargs == {}
+
+    minimal = adapter(ptr, value, mask)
+    assert minimal.args == (ptr, mask, None)
+    assert minimal.kwargs == {}
+
+
+def test_triton_load_adapter_keys_passthrough():
+    """Adapter mirrors tl.load(ptr, mask=..., other=..., keys=...)."""
+    adapter = TRITON_ADAPTERS[Load]
+    ptr, mask, other, keys = object(), object(), object(), object()
+
+    cache_modifier = "wb"
+    eviction_policy = "evict_last"
+    is_volatile = False
+
+    full = adapter(
+        ptr,
+        mask,
+        other,
+        cache_modifier=cache_modifier,
+        eviction_policy=eviction_policy,
+        is_volatile=is_volatile,
+        keys=keys,
+    )
+    assert full.args == (ptr, mask, keys)
+    assert full.kwargs == {}
+
+    minimal = adapter(ptr, mask, other)
+    assert minimal.args == (ptr, mask, None)
+    assert minimal.kwargs == {}
+
+
+def test_triton_reduce_adapter_supports_positional_and_keyword_axis():
+    """Adapter retains axis/keep_dims regardless of positional or keyword usage."""
+    adapter = TRITON_ADAPTERS[ReduceSum]
+    tensor = object()
+    axis = 1
+    keep_dims = True
+
+    positional = adapter(tensor, axis, keep_dims)
+    assert positional.args == (tensor, axis, keep_dims)
+    assert positional.kwargs == {}
+
+    keyword = adapter(tensor, axis=2, keep_dims=False)
+    assert keyword.args == (tensor, 2, False)
+    assert keyword.kwargs == {}
+
+
+def test_triton_addptr_adapter_orders_arguments():
+    """Adapter maps tl.addptr(ptr, offset) to (ptr, offset)."""
+    adapter = TRITON_ADAPTERS[AddPtr]
+    ptr = object()
+    offset = object()
+    result = adapter(ptr, offset)
+    assert result.args == (ptr, offset)
+    assert result.kwargs == {}
+
+
+def test_program_id_adapter_returns_axis_only():
+    """Adapter maps tl.program_id(axis) to (axis,)."""
+    adapter = TRITON_ADAPTERS[ProgramId]
+    program_id = 0
+    result = adapter(program_id)
+    assert result.args == (program_id,)
+    assert result.kwargs == {}
+
+
+def test_nki_load_store_adapters_align_with_clients():
+    """NKI adapters normalize masked load/store to (tensor, mask, keys)."""
+    src = object()
+    dst = object()
+    keys = object()
+    mask = object()
+    value = object()
+
+    load_result = NKI_ADAPTERS[Load](src, keys, mask=mask)
+    assert load_result.args == (src, mask, keys)
+    assert load_result.kwargs == {}
+
+    store_result = NKI_ADAPTERS[Store](dst, keys, value, mask=mask)
+    assert store_result.args == (dst, mask, keys)
+    assert store_result.kwargs == {}
+
+
+def test_patchop_uses_adapter_for_callbacks():
+    """PatchOp must feed the adapted arguments into before/after callbacks."""
+    before_log: list[tuple[str, tuple, dict]] = []
+    after_log: list[tuple[str, tuple, dict]] = []
+
+    def before_callback(*args, **kwargs):
+        before_log.append(("before", args, kwargs))
+
+    def after_callback(ret, *args, **kwargs):
+        after_log.append(("after", args, kwargs))
+
+    adapter = TRITON_ADAPTERS[Store]
+    callbacks = OpCallbacks(
+        before_callback=before_callback, after_callback=after_callback
+    )
+    patch_op = PatchOp(
+        op=lambda *_a, **_k: "return-value",
+        op_type=Store,
+        callbacks=callbacks,
+        adapter=adapter,
+    )
+
+    ptr = object()
+    mask = object()
+    patch_op(ptr, "value", mask)
+
+    expected_args = (ptr, mask, None)
+    assert before_log == [("before", expected_args, {})]
+    assert after_log == [("after", expected_args, {})]
diff --git a/triton_viz/core/patch.py b/triton_viz/core/patch.py
index c38a5b80..2efaa777 100644
--- a/triton_viz/core/patch.py
+++ b/triton_viz/core/patch.py
@@ -78,7 +78,7 @@ def __init__(self, *args: Any, **kwargs: Any) -> None:
 
 def passthrough_adapter(*args: Any, **kwargs: Any) -> AdapterResult:
     """Return arguments unchanged for clients that expect the original signature."""
-    return AdapterResult(args, kwargs)
+    return AdapterResult(*args, **kwargs)
 
 
 def _program_id_adapter(axis: Any, *_args: Any, **_kwargs: Any) -> AdapterResult:
@@ -105,8 +105,7 @@ def _triton_raw_load_adapter(ptr: Any, *_args: Any, **_kwargs: Any) -> AdapterRe
 def _triton_load_adapter(
     ptr: Any, mask: Any, _other: Any, *_args: Any, **kwargs: Any
 ) -> AdapterResult:
-    keys = kwargs.get("keys")
-    return AdapterResult(ptr, mask, keys)
+    return AdapterResult(ptr, mask, kwargs.get("keys"))
 
 
 def _triton_dot_adapter(a: Any, b: Any, *_args: Any, **_kwargs: Any) -> AdapterResult:

From 71b8e71070ad1b7458b5a9c05c5db39a494371f1 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Thu, 6 Nov 2025 14:47:43 -0500
Subject: [PATCH 75/90] revert unwanted changes from main

---
 triton_viz/clients/sanitizer/sanitizer.py |  4 ++--
 triton_viz/clients/tracer/tracer.py       | 12 ++++++------
 triton_viz/core/patch.py                  |  6 +++++-
 3 files changed, 13 insertions(+), 9 deletions(-)

diff --git a/triton_viz/clients/sanitizer/sanitizer.py b/triton_viz/clients/sanitizer/sanitizer.py
index d5c4edc0..8ea4dee7 100644
--- a/triton_viz/clients/sanitizer/sanitizer.py
+++ b/triton_viz/clients/sanitizer/sanitizer.py
@@ -548,7 +548,7 @@ def grid_callback(self, grid: tuple[int, ...]) -> None:
         self.last_grid = _get_last_grid(grid)
         self.tensors = sorted(self.tensors, key=lambda x: x.data_ptr())
 
-    def register_op_callback(self, op_type: type[Op], *args, **kwargs) -> OpCallbacks:
+    def register_op_callback(self, op_type: type[Op]) -> OpCallbacks:
         def pre_load_callback(ptr, mask, _keys):
             first_loc = np.unravel_index(np.argmax(mask, axis=None), mask.data.shape)
             first_ptr = ptr.data[first_loc]
@@ -1685,7 +1685,7 @@ def grid_callback(self, grid: tuple[int, ...]) -> None:
     def grid_idx_callback(self, grid_idx: tuple[int, ...]) -> None:
         self.grid_idx = grid_idx
 
-    def register_op_callback(self, op_type: type[Op], *args, **kwargs) -> OpCallbacks:
+    def register_op_callback(self, op_type: type[Op]) -> OpCallbacks:
         def op_program_id_overrider(axis):
             assert self.grid, "Grid not initialized!"
             return SymbolicExpr("pid", self.grid, axis)
diff --git a/triton_viz/clients/tracer/tracer.py b/triton_viz/clients/tracer/tracer.py
index af801221..bc7c1def 100644
--- a/triton_viz/clients/tracer/tracer.py
+++ b/triton_viz/clients/tracer/tracer.py
@@ -83,7 +83,7 @@ def grid_idx_callback(self, grid_idx: tuple[int, ...]):
     def grid_callback(self, grid: tuple[int, ...]):
         self.tensors = sorted(self.tensors, key=lambda x: x.data_ptr())
 
-    def register_op_callback(self, op_type: type[Op], *args, **kwargs) -> OpCallbacks:
+    def register_op_callback(self, op_type: type[Op]) -> OpCallbacks:
         def _extract_user_frames() -> list[traceback.FrameSummary]:
             stack: list[traceback.FrameSummary] = list(traceback.extract_stack())
             # drop current frames (this function and callers)
@@ -126,7 +126,7 @@ def _convert_keys_to_numpy(keys):
             else:
                 return keys
 
-        def pre_masked_load_callback(ptr, mask, keys):
+        def pre_load_callback(ptr, mask, keys):
             if not self.sample:
                 return
 
@@ -143,7 +143,7 @@ def pre_masked_load_callback(ptr, mask, keys):
             rec.call_path = _extract_user_frames()
             self.records.append(rec)
 
-        def pre_masked_store_callback(ptr, mask, keys):
+        def pre_store_callback(ptr, mask, keys):
             if not self.sample:
                 return
 
@@ -189,7 +189,7 @@ def pre_raw_store_callback(ptr, value):
             rec.call_path = _extract_user_frames()
             self.records.append(rec)
 
-        def post_reduce_sum_callback(ret, input, axis, keep_dims):
+        def post_reduce_sum_callback(ret, input, axis=None, keep_dims=False):
             if not self.sample:
                 return
             input_shape = input.handle.data.shape
@@ -231,9 +231,9 @@ def post_flip_callback(ret, x, *args, **kwargs):
         if op_type is Allocate:
             return OpCallbacks(after_callback=post_allocate_callback)
         elif op_type is Load:
-            return OpCallbacks(before_callback=pre_masked_load_callback)
+            return OpCallbacks(before_callback=pre_load_callback)
         elif op_type is Store:
-            return OpCallbacks(before_callback=pre_masked_store_callback)
+            return OpCallbacks(before_callback=pre_store_callback)
         elif op_type is RawLoad:
             return OpCallbacks(before_callback=pre_raw_load_callback)
         elif op_type is RawStore:
diff --git a/triton_viz/core/patch.py b/triton_viz/core/patch.py
index c3d6c261..c3709f40 100644
--- a/triton_viz/core/patch.py
+++ b/triton_viz/core/patch.py
@@ -13,8 +13,8 @@
     Allocate,
     RawLoad,
     Load,
-    Store,
     RawStore,
+    Store,
     UnaryOp,
     BinaryOp,
     TernaryOp,
@@ -45,6 +45,7 @@
     CumSum,
     Bitcast,
     AtomicCas,
+    AtomicRMW,
 )
 from .data import Flip  # separate import to avoid reordering noise
 import inspect
@@ -179,6 +180,7 @@ def _nki_dot_adapter(x: Any, y: Any, *_args: Any, **_kwargs: Any) -> AdapterResu
     CumSum,
     Bitcast,
     AtomicCas,
+    AtomicRMW,
 ]
 
 TRITON_ORIGINAL_OPS = {
@@ -212,6 +214,7 @@ def _nki_dot_adapter(x: Any, y: Any, *_args: Any, **_kwargs: Any) -> AdapterResu
     Trans: interpreter_builder.create_trans,
     Bitcast: interpreter_builder.create_bitcast,
     AtomicCas: interpreter_builder.create_atomic_cas,
+    AtomicRMW: interpreter_builder.create_atomic_rmw,
 }
 
 TRITON_OP_ATTR_NAMES = {
@@ -245,6 +248,7 @@ def _nki_dot_adapter(x: Any, y: Any, *_args: Any, **_kwargs: Any) -> AdapterResu
     Trans: "create_trans",
     Bitcast: "create_bitcast",
     AtomicCas: "create_atomic_cas",
+    AtomicRMW: "create_atomic_rmw",
 }
 
 TRITON_ADAPTERS: dict[type[Op], Callable[..., AdapterResult]] = {

From cc9a787875d8f02239abf287f79933822fb230a2 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Thu, 6 Nov 2025 14:48:32 -0500
Subject: [PATCH 76/90] patch backend less bad

---
 triton_viz/core/client.py |  2 +-
 triton_viz/core/patch.py  | 66 +++++++++++----------------------------
 2 files changed, 19 insertions(+), 49 deletions(-)

diff --git a/triton_viz/core/client.py b/triton_viz/core/client.py
index a6b8bc3b..254c9db4 100644
--- a/triton_viz/core/client.py
+++ b/triton_viz/core/client.py
@@ -148,7 +148,7 @@ def patch_run(self, fn, backend: str):
                 backend_ops = OPERATION_REGISTRY[backend]["op_list"]
 
                 for op in backend_ops:
-                    unpatch_op(op)
+                    unpatch_op(op, backend)
                 unpatch_for_loop()
                 unpatch_lang(backend)
 
diff --git a/triton_viz/core/patch.py b/triton_viz/core/patch.py
index c3709f40..1358b3d8 100644
--- a/triton_viz/core/patch.py
+++ b/triton_viz/core/patch.py
@@ -310,12 +310,14 @@ def _nki_dot_adapter(x: Any, y: Any, *_args: Any, **_kwargs: Any) -> AdapterResu
 
 OPERATION_REGISTRY: dict[str, dict[str, Any]] = {
     "triton": {
+        "builder": interpreter_builder,
         "op_list": TRITON_OP_LIST,
         "original_ops": TRITON_ORIGINAL_OPS,
         "op_attr_names": TRITON_OP_ATTR_NAMES,
         "adapters": TRITON_ADAPTERS,
     },
     "nki": {
+        "builder": nki_builder,
         "op_list": NKI_OP_LIST,
         "original_ops": NKI_ORIGINAL_OPS,
         "op_attr_names": NKI_OP_ATTR_NAMES,
@@ -323,34 +325,6 @@ def _nki_dot_adapter(x: Any, y: Any, *_args: Any, **_kwargs: Any) -> AdapterResu
     },
 }
 
-BUILDER = interpreter_builder
-current_backend = "triton"
-op_list: list[Op] = OPERATION_REGISTRY[current_backend]["op_list"]
-original_ops: dict[Op, Callable] = OPERATION_REGISTRY[current_backend]["original_ops"]
-_OP_ATTR_NAMES: dict[Op, str] = OPERATION_REGISTRY[current_backend]["op_attr_names"]
-
-
-def get_builder_for_backend(backend: str):
-    """Get the appropriate builder for a given backend."""
-    return nki_builder if backend == "nki" else interpreter_builder
-
-
-def detect_current_backend():
-    """Detect the current backend based on which operations have been patched."""
-    # Check if Triton operations have been patched first (since this example uses Triton)
-    if hasattr(interpreter_builder, "create_get_program_id") and hasattr(
-        interpreter_builder.create_get_program_id, "__wrapped__"
-    ):
-        return "triton"
-    # Check if NKI operations have been patched
-    elif hasattr(nki_builder, "program_id") and hasattr(
-        nki_builder.program_id, "__wrapped__"
-    ):
-        return "nki"
-    # Default to checking BUILDER global
-    else:
-        return "nki" if BUILDER == nki_builder else "triton"
-
 
 reduce_map: dict[type[Op], Callable] = {
     ReduceMax: tl.max,
@@ -408,7 +382,7 @@ def __call__(self, *args, **kwargs):
         return ret
 
 
-def patch_op(op_type: type[Op], callbacks: OpCallbacks, backend: Optional[str] = None):
+def patch_op(op_type: type[Op], callbacks: OpCallbacks, backend: str):
     """
     Register a callback to be called before and after an operator is executed.
 
@@ -416,27 +390,20 @@ def patch_op(op_type: type[Op], callbacks: OpCallbacks, backend: Optional[str] =
     :param callbacks: The OpCallbacks object containing before_callback, after_callback, and op_overrider.
     :param backend: The backend to use ('triton', 'nki', or None for current backend).
     """
-    if backend is None:
-        backend = current_backend
-
     if backend not in OPERATION_REGISTRY:
         raise ValueError(f"Unknown backend: {backend}")
 
     backend_ops = OPERATION_REGISTRY[backend]["original_ops"]
     backend_attr_names = OPERATION_REGISTRY[backend]["op_attr_names"]
     backend_adapters = OPERATION_REGISTRY[backend]["adapters"]
-    backend_builder = nki_builder if backend == "nki" else interpreter_builder
+    backend_builder = OPERATION_REGISTRY[backend]["builder"]
 
     if op_type in backend_ops:
         op_name = backend_attr_names[op_type]
         original_op = backend_ops[op_type]
         adapter = backend_adapters[op_type]
         patched_op = PatchOp(original_op, op_type, callbacks, adapter)
-        setattr(
-            backend_builder,
-            op_name,
-            lambda *args, **kwargs: patched_op(*args, **kwargs),
-        )
+        setattr(backend_builder, op_name, patched_op)
     elif backend == "triton" and (op_type in reduce_map or op_type in scan_map):
         if op_type in reduce_map:
             op_name = reduce_map[op_type].__name__
@@ -445,28 +412,32 @@ def patch_op(op_type: type[Op], callbacks: OpCallbacks, backend: Optional[str] =
         original_op = getattr(tl, op_name)
         adapter = backend_adapters[op_type]
         patched_op = PatchOp(original_op, op_type, callbacks, adapter)
-        setattr(tl, op_name, lambda *args, **kwargs: patched_op(*args, **kwargs))
+        setattr(tl, op_name, patched_op)
     elif backend == "triton" and op_type in math_map:
         op_name = math_map[op_type].__name__
         original_op = getattr(tl.math, op_name)
         adapter = backend_adapters[op_type]
         patched_op = PatchOp(original_op, op_type, callbacks, adapter)
-        setattr(tl.math, op_name, lambda *args, **kwargs: patched_op(*args, **kwargs))
+        setattr(tl.math, op_name, patched_op)
     else:
         raise ValueError(f"Patching operator {op_type} not supported")
 
 
-def unpatch_op(op_type: type[Op]):
+def unpatch_op(op_type: type[Op], backend: str):
     """
     Unregister a callback for an operator.
 
     :param op_type: The type of the operator to unregister the callback for.
     """
-    if op_type in original_ops:
-        original_op = original_ops[op_type]  # type: ignore
+    backend_ops = OPERATION_REGISTRY[backend]["original_ops"]
+    backend_attr_names = OPERATION_REGISTRY[backend]["op_attr_names"]
+    backend_builder = OPERATION_REGISTRY[backend]["builder"]
+
+    if op_type in backend_ops:
+        original_op = backend_ops[op_type]  # type: ignore
         # Use hardcoded name from _OP_ATTR_NAMES
-        op_name = _OP_ATTR_NAMES[op_type]  # type: ignore
-        setattr(BUILDER, op_name, original_op)
+        op_name = backend_attr_names[op_type]  # type: ignore
+        setattr(backend_builder, op_name, original_op)
 
 
 class _LoopIter:
@@ -794,6 +765,8 @@ def _grid_executor_call(self, *args_dev, backend=None, **kwargs):
     if kwargs.pop("warmup", False):
         return
 
+    builder = OPERATION_REGISTRY[backend]["builder"]
+
     def run_grid_loops(grid):
         for x in tqdm(
             range(grid[0]),
@@ -853,9 +826,6 @@ def run_grid_loops(grid):
     assert len(grid) <= 3
     grid = grid + (1,) * (3 - len(grid))
 
-    # Use the correct builder based on current backend
-    builder = get_builder_for_backend(backend)
-
     if backend == "nki":
         builder.set_grid_dim(grid)
     else:  # triton

From f94d1be51533ff0ef0eb75576ef736f8047ece65 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Thu, 6 Nov 2025 14:49:10 -0500
Subject: [PATCH 77/90] remove unhelpful comments + show expected transformed
 code

---
 triton_viz/core/nki_extract_slice.py | 73 +++++++++++-----------------
 1 file changed, 29 insertions(+), 44 deletions(-)

diff --git a/triton_viz/core/nki_extract_slice.py b/triton_viz/core/nki_extract_slice.py
index 7b1c5267..5b7b9188 100644
--- a/triton_viz/core/nki_extract_slice.py
+++ b/triton_viz/core/nki_extract_slice.py
@@ -1,5 +1,6 @@
 import ast
 
+
 class StoreCallTransformer(ast.NodeTransformer):
     """
     A targeted AST transformer to rewrite `nl.store(x[...], ...)` calls
@@ -7,7 +8,6 @@ class StoreCallTransformer(ast.NodeTransformer):
     into `masked_load(x, slice_obj)` with a preceding assignment.
     """
 
-    # MODIFICATION 1: We now visit the 'Expr' statement wrapper.
     def visit_Expr(self, node: ast.Expr) -> ast.AST | list[ast.AST]:
         """
         Intercept and transform expression statements.
@@ -18,100 +18,83 @@ def visit_Expr(self, node: ast.Expr) -> ast.AST | list[ast.AST]:
 
         call_node = node.value
 
-        # --- The logic from the old visit_Call starts here ---
-
-        # 1. IDENTIFY: Is this the specific call we want to transform?
         if not (
             isinstance(call_node.func, ast.Attribute)
             and isinstance(call_node.func.value, ast.Name)
-            and call_node.func.value.id == 'nl'
-            and call_node.func.attr in ('store', 'load')
+            and call_node.func.value.id == "nl"
+            and call_node.func.attr in ("store", "load")
         ):
             return self.generic_visit(node)
 
-        # 2. VALIDATE: Does the call have the expected structure?
         if not call_node.args or not isinstance(call_node.args[0], ast.Subscript):
             return self.generic_visit(node)
 
-        # 3. DECONSTRUCT
         subscript_node = call_node.args[0]
         sliced_object = subscript_node.value
         slice_content = subscript_node.slice
         remaining_args = call_node.args[1:]
 
-        # 4. REBUILD
         # Convert to nl.masked_load or nl.masked_store
-        func_name = 'masked_' + call_node.func.attr
+        func_name = "masked_" + call_node.func.attr
 
         new_call = ast.Call(
             func=ast.Attribute(
-                value=ast.Name(id='nl', ctx=ast.Load()),
-                attr=func_name,
-                ctx=ast.Load()
+                value=ast.Name(id="nl", ctx=ast.Load()), attr=func_name, ctx=ast.Load()
             ),
             args=[
                 sliced_object,
                 self._create_slice_value_node(slice_content),
-                *remaining_args
+                *remaining_args,
             ],
-            keywords=call_node.keywords
+            keywords=call_node.keywords,
         )
 
-        # 5. REPLACE
         return ast.Expr(value=new_call)
 
     def visit_Assign(self, node: ast.Assign) -> ast.AST:
         """
         Handle assignment statements where the value is an nl.load or nl.store call.
         """
-        # Check if the value being assigned is a function call
         if not isinstance(node.value, ast.Call):
             return self.generic_visit(node)
 
         call_node = node.value
 
-        # 1. IDENTIFY: Is this the specific call we want to transform?
         if not (
             isinstance(call_node.func, ast.Attribute)
             and isinstance(call_node.func.value, ast.Name)
-            and call_node.func.value.id == 'nl'
-            and call_node.func.attr in ('store', 'load')
+            and call_node.func.value.id == "nl"
+            and call_node.func.attr in ("store", "load")
         ):
             return self.generic_visit(node)
 
-        # 2. VALIDATE: Does the call have the expected structure?
         if not call_node.args or not isinstance(call_node.args[0], ast.Subscript):
             return self.generic_visit(node)
 
-        # 3. DECONSTRUCT
         subscript_node = call_node.args[0]
         sliced_object = subscript_node.value
         slice_content = subscript_node.slice
         remaining_args = call_node.args[1:]
 
-        # 4. REBUILD
         # Convert to nl.masked_load or nl.masked_store
-        func_name = 'masked_' + call_node.func.attr
+        func_name = "masked_" + call_node.func.attr
 
         new_call = ast.Call(
             func=ast.Attribute(
-                value=ast.Name(id='nl', ctx=ast.Load()),
-                attr=func_name,
-                ctx=ast.Load()
+                value=ast.Name(id="nl", ctx=ast.Load()), attr=func_name, ctx=ast.Load()
             ),
             args=[
                 sliced_object,
                 self._create_slice_value_node(slice_content),
-                *remaining_args
+                *remaining_args,
             ],
-            keywords=call_node.keywords
+            keywords=call_node.keywords,
         )
 
-        # 5. REPLACE: Create a new assignment with the transformed call
         return ast.Assign(
             targets=node.targets,
             value=new_call,
-            type_comment=getattr(node, 'type_comment', None)
+            type_comment=getattr(node, "type_comment", None),
         )
 
     def _create_slice_value_node(self, node: ast.AST) -> ast.AST:
@@ -122,22 +105,23 @@ def _create_slice_value_node(self, node: ast.AST) -> ast.AST:
         match node:
             case ast.Slice(lower, upper, step):
                 return ast.Call(
-                    func=ast.Name(id='slice', ctx=ast.Load()),
+                    func=ast.Name(id="slice", ctx=ast.Load()),
                     args=[
                         lower or ast.Constant(value=None),
                         upper or ast.Constant(value=None),
-                        step or ast.Constant(value=None)
+                        step or ast.Constant(value=None),
                     ],
-                    keywords=[]
+                    keywords=[],
                 )
             case ast.Tuple(elts):
                 return ast.Tuple(
                     elts=[self._create_slice_value_node(e) for e in elts],
-                    ctx=ast.Load()
+                    ctx=ast.Load(),
                 )
             case _:
                 return node
 
+
 def transform_code(source_code: str) -> str:
     """
     Applies the StoreCallTransformer to a string of Python code.
@@ -148,7 +132,7 @@ def transform_code(source_code: str) -> str:
     ast.fix_missing_locations(new_tree)
     return ast.unparse(new_tree)
 
-# --- Input Code ---
+
 source_code = """
 import numpy as np
 
@@ -166,12 +150,13 @@ def transform_code(source_code: str) -> str:
 nl.some_other_func(z[:5])
 """
 
-if __name__ == '__main__':
-    # --- Transformation ---
-    transformed_code = transform_code(source_code)
+expected_transformed_code = """\
+import numpy as np
+sbuf_value = nl.masked_load(x, (slice(2, 4, None), None, ..., nl.arange(128)[None, :]), mask=mask)
+nl.masked_store(x, (slice(2, 4, None), None, ..., nl.arange(128)[None, :]), sbuf_value, mask=mask)
+other_lib.store(y[10:])
+nl.some_other_func(z[:5])\
+"""
 
-    # --- Output ---
-    print("--- Original Code ---")
-    print(source_code)
-    print("\n--- Transformed Code ---")
-    print(transformed_code)
+if __name__ == "__main__":
+    assert transform_code(source_code) == expected_transformed_code

From 675ababfca5cf0e2fbf6856c2f17cf422bb01225 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Thu, 6 Nov 2025 14:49:39 -0500
Subject: [PATCH 78/90] debloat nki offsets code

---
 triton_viz/core/nki.py | 39 +++++----------------------------------
 1 file changed, 5 insertions(+), 34 deletions(-)

diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index c0860cd9..5ee7bec4 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -47,42 +47,13 @@ def cpu(self):  # THTODO: rm?
     def get_offsets(self):
         """
         Generate offset arrays for each dimension based on shape and stride.
-
-        Args:
-            strides: Tuple of strides for each dimension (a, b, ..., z)
-
-        Returns:
-            Tuple of offset arrays: (arange(A)[:, None, ..., None]*a, arange(B)[None, :, ..., None]*b, ...)
+        Given array with shape (A, ..., Z) and strides (a, ..., z), return offsets:
+        a * arange(A)[:, None, ..., None] + ... + z * arange(Z)[None, None, ..., :]
         """
-        strides = self.data.strides
-        if self.data is None:
-            raise AttributeError("NDArray has no value - cannot compute offsets")
-
-        shape = self.shape
-        if len(shape) != len(strides):
-            raise ValueError(
-                f"Shape has {len(shape)} dimensions but strides has {len(strides)} dimensions"
-            )
-
-        # offsets = []
         offsets = 0
-        ndim = len(shape)
-
-        for i, (dim_size, stride) in enumerate(zip(shape, strides)):
-            # Create arange for this dimension
-            arange_vals = np.arange(dim_size)
-
-            # Create broadcast shape - put arange in position i, others as 1
-            broadcast_shape = [1] * ndim
-            broadcast_shape[i] = dim_size
-
-            # Reshape and multiply by stride
-            offset_array = (arange_vals * stride).reshape(broadcast_shape)
-            # offsets.append(NDArray(value=offset_array, name=f"{self.name}_offset_dim{i}"))
-            offsets += NDArray(value=offset_array, name=f"{self.name}_offset_dim{i}")
-
-        # return tuple(offsets)
-        return offsets
+        for dim_size, stride in zip(self.shape, self.stride()):
+            offsets = np.expand_dims(offsets, -1) + np.arange(dim_size) * stride
+        return NDArray(value=offsets, name=self.name)
 
     def __repr__(self):
         return f"NDArray(shape={self.shape}, dtype={self.dtype}, name={self.name})"

From 7c3bec2262a236e165cd419ee5ce3dff3a592129 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Thu, 6 Nov 2025 14:50:13 -0500
Subject: [PATCH 79/90] fix matmul visualization error

---
 triton_viz/core/nki.py             | 8 +++++++-
 triton_viz/visualizer/interface.py | 5 +----
 2 files changed, 8 insertions(+), 5 deletions(-)

diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index 5ee7bec4..01d3c3bf 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -41,9 +41,15 @@ def stride(self):
     def element_size(self):
         return self.dtype.itemsize
 
-    def cpu(self):  # THTODO: rm?
+    def cpu(self):
         return self
 
+    def detach(self):
+        return self
+
+    def numpy(self):
+        return self.data
+
     def get_offsets(self):
         """
         Generate offset arrays for each dimension based on shape and stride.
diff --git a/triton_viz/visualizer/interface.py b/triton_viz/visualizer/interface.py
index f52cdb42..00c9947d 100644
--- a/triton_viz/visualizer/interface.py
+++ b/triton_viz/visualizer/interface.py
@@ -3,7 +3,6 @@
 from .analysis import analyze_records
 from .draw import get_visualization_data
 import os
-import torch
 from flask_cloudflared import _run_cloudflared
 import requests
 import time
@@ -59,9 +58,7 @@ def precompute_c_values(op_data):
         for j in range(cols):
             precomputed[(i, j)] = [0] * (inner_dim + 1)
             for k in range(1, inner_dim + 1):
-                precomputed[(i, j)][k] = torch.dot(
-                    input_data[i, :k], other_data[:k, j]
-                ).item()
+                precomputed[(i, j)][k] = input_data[i, :k] @ other_data[:k, j]
 
     return precomputed
 

From 229d4a9e4e5caeb2bf6a6fe11af0de84afe9642b Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Thu, 6 Nov 2025 14:55:08 -0500
Subject: [PATCH 80/90] lint

---
 triton_viz/static/flip_3d.js | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/triton_viz/static/flip_3d.js b/triton_viz/static/flip_3d.js
index c8351952..f87f7fee 100644
--- a/triton_viz/static/flip_3d.js
+++ b/triton_viz/static/flip_3d.js
@@ -96,34 +96,34 @@ export function createFlip3D(containerElement, options) {
     const rowGap = CUBE_SIZE * 4.2; // vertical distance per step row
     const columnGap = CUBE_SIZE * 6.0; // horizontal gap between view and swap panels
 
-    let currentY = 0;
+    let current_y = 0;
     // initial 1D row (single panel)
     const initPanel = new THREE.Group();
     for (let i=0;i<length;i++){
         const color = hsv(hueAt(i), 1, 1);
         const cube = createCube(color, 'Flip', i, 0, 0, cubeGeometry, edgesGeometry, lineMaterial);
-        cube.position.set(i*(CUBE_SIZE+GAP), currentY, 0);
+        cube.position.set(i*(CUBE_SIZE+GAP), current_y, 0);
         initPanel.add(cube);
     }
     const initLabel = createTextSprite('initial array');
-    initLabel.position.set(-CUBE_SIZE*3.2, currentY + CUBE_SIZE*0.5, 0);
+    initLabel.position.set(-CUBE_SIZE*3.2, current_y + CUBE_SIZE*0.5, 0);
     initPanel.add(initLabel);
     root.add(initPanel); layers.push(initPanel);
 
-    currentY -= rowGap;
+    current_y -= rowGap;
     let arr = Array.from({length}, (_,i)=>i);
     for (let s=0; s<steps.length; s++){
         const seg = steps[s];
         const lw = layoutWidth(seg, length);
         const viewX = 0;
         const swapX = viewX + lw + columnGap;
-        const viewPanel = buildPanel(arr, seg, currentY, `step ${s} view`, viewX);
+        const viewPanel = buildPanel(arr, seg, current_y, `step ${s} view`, viewX);
         const swapped = swapBlocksArray(arr, seg);
-        const swapPanel = buildPanel(swapped, seg, currentY, `step ${s} swap`, swapX);
+        const swapPanel = buildPanel(swapped, seg, current_y, `step ${s} swap`, swapX);
         root.add(viewPanel); root.add(swapPanel);
         layers.push(viewPanel); layers.push(swapPanel);
         arr = swapped.slice();
-        currentY -= rowGap;
+        current_y -= rowGap;
     }
 
     // final reshape (1D row)
@@ -131,11 +131,11 @@ export function createFlip3D(containerElement, options) {
     for (let i=0;i<arr.length;i++){
         const color = hsv(hueAt(arr[i]), 1, 1);
         const cube = createCube(color, 'Flip', i, 0, 0, cubeGeometry, edgesGeometry, lineMaterial);
-        cube.position.set(i*(CUBE_SIZE+GAP), currentY, 0);
+        cube.position.set(i*(CUBE_SIZE+GAP), current_y, 0);
         finalPanel.add(cube);
     }
     const finalLabel = createTextSprite('final reshape');
-    finalLabel.position.set(-CUBE_SIZE*3.2, currentY + CUBE_SIZE*0.5, 0);
+    finalLabel.position.set(-CUBE_SIZE*3.2, current_y + CUBE_SIZE*0.5, 0);
     finalPanel.add(finalLabel);
     root.add(finalPanel); layers.push(finalPanel);
 

From 54559f1b7764a2c77a33985064a016dda3b9f12a Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Thu, 13 Nov 2025 10:34:48 -0500
Subject: [PATCH 81/90] rename nki_masked_load -> masked_load since it's also
 used for triton

---
 tests/{nki => }/test_masked_load.py                    | 2 +-
 triton_viz/clients/tracer/tracer.py                    | 2 +-
 triton_viz/core/__init__.py                            | 2 +-
 triton_viz/core/{nki_masked_load.py => masked_load.py} | 0
 triton_viz/core/nki.py                                 | 2 +-
 5 files changed, 4 insertions(+), 4 deletions(-)
 rename tests/{nki => }/test_masked_load.py (99%)
 rename triton_viz/core/{nki_masked_load.py => masked_load.py} (100%)

diff --git a/tests/nki/test_masked_load.py b/tests/test_masked_load.py
similarity index 99%
rename from tests/nki/test_masked_load.py
rename to tests/test_masked_load.py
index 1c0ebf3b..f966d808 100644
--- a/tests/nki/test_masked_load.py
+++ b/tests/test_masked_load.py
@@ -1,6 +1,6 @@
 import numpy as np
 import pytest
-from triton_viz.core.nki_masked_load import masked_load, masked_store
+from triton_viz.core.masked_load import masked_load, masked_store
 
 
 def print_op_details(
diff --git a/triton_viz/clients/tracer/tracer.py b/triton_viz/clients/tracer/tracer.py
index bc7c1def..b9faf201 100644
--- a/triton_viz/clients/tracer/tracer.py
+++ b/triton_viz/clients/tracer/tracer.py
@@ -12,7 +12,7 @@
     Allocate,
     Flip,
 )
-from triton_viz.core.nki_masked_load import masked_load
+from triton_viz.core.masked_load import masked_load
 from typing import Callable, Optional, Union
 import numpy as np
 import traceback
diff --git a/triton_viz/core/__init__.py b/triton_viz/core/__init__.py
index 1f66dd83..73512549 100644
--- a/triton_viz/core/__init__.py
+++ b/triton_viz/core/__init__.py
@@ -26,7 +26,7 @@
     Rsqrt,
     CastImpl,
 )
-from .nki_masked_load import masked_load, masked_store
+from .masked_load import masked_load, masked_store
 from .nki_extract_slice import StoreCallTransformer, transform_code
 
 __all__ = [
diff --git a/triton_viz/core/nki_masked_load.py b/triton_viz/core/masked_load.py
similarity index 100%
rename from triton_viz/core/nki_masked_load.py
rename to triton_viz/core/masked_load.py
diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index 01d3c3bf..d7726fcb 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -3,7 +3,7 @@
 import neuronxcc.nki.language as nl
 import inspect
 from .nki_extract_slice import transform_code
-from .nki_masked_load import masked_load, masked_store
+from .masked_load import masked_load, masked_store
 
 
 class NDArray:

From f9c72e368ee41f943a275cbdf07b1a4a23d135fd Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Thu, 13 Nov 2025 10:39:57 -0500
Subject: [PATCH 82/90] make AWS neuron an optional feature of triton viz (and
 include CI)

---
 .github/workflows/python-app.yml | 47 ++++++++++++++++
 README.md                        | 17 ++++++
 pyproject.toml                   |  7 +++
 pytest.ini                       |  4 ++
 tests/nki/test_nki.py            |  3 +-
 tests/test_adapters.py           |  4 ++
 triton_viz/core/nki.py           |  9 +++-
 triton_viz/core/patch.py         | 92 ++++++++++++++++++--------------
 8 files changed, 142 insertions(+), 41 deletions(-)

diff --git a/.github/workflows/python-app.yml b/.github/workflows/python-app.yml
index bb729408..d77ec62f 100644
--- a/.github/workflows/python-app.yml
+++ b/.github/workflows/python-app.yml
@@ -60,3 +60,50 @@ jobs:
       run: |
         cd triton_viz
         python -m pytest tests
+
+  nki-tests:
+    name: AWS Neuron / NKI
+    runs-on: ubuntu-22.04
+    needs: build
+    env:
+      TRITON_INTERPRET: "1"
+      TRITON_VERSION: "3.4.0"
+      PIP_EXTRA_INDEX_URL: https://pip.repos.neuron.amazonaws.com
+
+    steps:
+    - uses: actions/checkout@v3
+      with:
+        path: 'triton_viz'
+
+    - name: Set up Python 3.10
+      uses: actions/setup-python@v3
+      with:
+        python-version: '3.10'
+
+    - name: Configure AWS Neuron apt repository
+      run: |
+        sudo tee /etc/apt/sources.list.d/neuron.list <<'EOF'
+        deb https://apt.repos.neuron.amazonaws.com $(lsb_release -cs) main
+        EOF
+        wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AWS-NEURON.PUB | sudo apt-key add -
+        sudo apt-get update
+        sudo apt-get install -y aws-neuronx-runtime-lib aws-neuronx-collectives aws-neuronx-tools
+
+    - name: Install AWS Neuron Python packages
+      run: |
+        python -m pip install --upgrade pip
+        python -m pip install --upgrade neuronx-cc==2.* libneuronxla==2.* torch-neuronx==2.* torchvision --extra-index-url https://pip.repos.neuron.amazonaws.com
+
+    - name: Install Triton
+      run: |
+        python -m pip install triton==${{ env.TRITON_VERSION }}
+
+    - name: Install Triton-Viz with NKI extras
+      run: |
+        cd triton_viz
+        python -m pip install -e .[nki]
+
+    - name: Run NKI pytest suite
+      run: |
+        cd triton_viz
+        python -m pytest tests/nki -m nki
diff --git a/README.md b/README.md
index 9f902b23..372b1355 100644
--- a/README.md
+++ b/README.md
@@ -71,6 +71,23 @@ pip install -e .
 
 You're all set!
 
+### Optional: Enable NKI Support
+
+If you want to exercise the Neuron Kernel Interface (NKI) interpreter or run the NKI-specific tests:
+
+1. Follow the [AWS Neuron Torch-NeuronX Ubuntu 22.04 setup guide](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/setup/neuron-setup/pytorch/neuronx/ubuntu/torch-neuronx-ubuntu22.html#setup-torch-neuronx-ubuntu22) to add the Neuron APT repository and install the required system packages (for example `aws-neuronx-tools`, `aws-neuronx-runtime-lib`, `aws-neuronx-collectives`, and their dependencies).
+2. Install Triton-Viz with the optional NKI extras so the Neuron Python packages (`neuronx-cc`, `libneuronxla`, `torch-neuronx`) are available:
+
+   ```sh
+   pip install -e .[nki]
+   # or pip install triton-viz[nki]
+   ```
+
+### Testing
+* To run core Triton-viz tests, run `pytest tests/`.
+* (if NKI installed) To run NKI-specific tests, run `pytest -m nki tests/`.
+* To run all tests (Triton + NKI), run `pytest -m "" tests/`.
+
 ## Working with Examples
 
 ```sh
diff --git a/pyproject.toml b/pyproject.toml
index 0442001c..8680eb04 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -46,3 +46,10 @@ homepage = "https://github.com/Deep-Learning-Profiling-Tools/triton-viz"
 [project.scripts]
 triton-sanitizer = "triton_viz.wrapper:apply_sanitizer"
 triton-profiler = "triton_viz.wrapper:apply_profiler"
+
+[project.optional-dependencies]
+nki = [
+  "neuronx-cc==2.*",
+  "libneuronxla==2.*",
+  "torch-neuronx==2.*",
+]
diff --git a/pytest.ini b/pytest.ini
index 0f5ccb8f..1d3f80a6 100644
--- a/pytest.ini
+++ b/pytest.ini
@@ -2,3 +2,7 @@
 python_files = *.py
 python_classes = *Test
 python_functions = test_*
+addopts = -m "not nki"
+
+markers =
+  nki: NKI-specific tests
diff --git a/tests/nki/test_nki.py b/tests/nki/test_nki.py
index 6c866ed0..01dc98cc 100644
--- a/tests/nki/test_nki.py
+++ b/tests/nki/test_nki.py
@@ -4,8 +4,9 @@
 """
 import numpy as np
 from triton_viz.core.nki import NDArray
+import pytest
 
-# sys.path.insert(0, os.path.join(os.path.dirname(__file__), "triton_viz", "core"))
+pytestmark = pytest.mark.nki  # only run at "pytest -m nki"
 
 
 def test_ndarray_creation():
diff --git a/tests/test_adapters.py b/tests/test_adapters.py
index f1345cb0..b22e24a1 100644
--- a/tests/test_adapters.py
+++ b/tests/test_adapters.py
@@ -1,3 +1,5 @@
+import pytest
+
 from triton_viz.core.callbacks import OpCallbacks
 from triton_viz.core.data import (
     AddPtr,
@@ -8,6 +10,7 @@
 )
 from triton_viz.core.patch import (
     AdapterResult,
+    HAS_NKI,
     NKI_ADAPTERS,
     PatchOp,
     TRITON_ADAPTERS,
@@ -120,6 +123,7 @@ def test_program_id_adapter_returns_axis_only():
     assert result.kwargs == {}
 
 
+@pytest.mark.skipif(not HAS_NKI, reason="NKI extras not installed")
 def test_nki_load_store_adapters_align_with_clients():
     """NKI adapters normalize masked load/store to (tensor, mask, keys)."""
     src = object()
diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index d7726fcb..9d7d97b2 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -1,6 +1,13 @@
 import numpy as np
 
-import neuronxcc.nki.language as nl
+try:
+    import neuronxcc.nki.language as nl
+except (
+    ModuleNotFoundError
+) as exc:  # pragma: no cover - only hit when optional deps missing
+    raise ModuleNotFoundError(
+        "NeuronX dependencies are missing. Install triton-viz[nki] to enable the NKI interpreter."
+    ) from exc
 import inspect
 from .nki_extract_slice import transform_code
 from .masked_load import masked_load, masked_store
diff --git a/triton_viz/core/patch.py b/triton_viz/core/patch.py
index 1358b3d8..f171c6f6 100644
--- a/triton_viz/core/patch.py
+++ b/triton_viz/core/patch.py
@@ -58,7 +58,15 @@
 from triton.runtime.interpreter import _patch_lang as triton_patch_lang
 from triton.runtime.interpreter import ASTTransformer as _OrigASTTransformer
 from triton.runtime import JITFunction
-from triton_viz.core.nki import nki_builder
+
+HAS_NKI = False
+nki_builder = None
+try:
+    from triton_viz.core.nki import nki_builder  # type: ignore
+
+    HAS_NKI = True
+except ModuleNotFoundError:
+    pass
 
 
 @dataclass
@@ -265,47 +273,53 @@ def _nki_dot_adapter(x: Any, y: Any, *_args: Any, **_kwargs: Any) -> AdapterResu
 for op_type in TRITON_OP_LIST:
     TRITON_ADAPTERS.setdefault(op_type, passthrough_adapter)
 
+NKI_OP_LIST: list[type[Op]] = []
+NKI_ORIGINAL_OPS: dict[type[Op], Callable] = {}
+NKI_OP_ATTR_NAMES: dict[type[Op], str] = {}
+NKI_ADAPTERS: dict[type[Op], Callable[..., AdapterResult]] = {}
+if HAS_NKI:
+    assert nki_builder is not None
+
+    NKI_OP_LIST = [
+        Allocate,
+        ProgramId,
+        Load,
+        Store,
+        Dot,
+        UnaryOp,
+        MakeRange,
+    ]
+
+    NKI_ORIGINAL_OPS = {
+        ProgramId: nki_builder.program_id,
+        Allocate: nki_builder.ndarray,
+        Load: nki_builder.masked_load,
+        Store: nki_builder.masked_store,
+        Dot: nki_builder.matmul,
+        UnaryOp: nki_builder._unary_op,
+        MakeRange: nki_builder.arange,
+    }
 
-NKI_OP_LIST = [
-    Allocate,
-    ProgramId,
-    Load,
-    Store,
-    Dot,
-    UnaryOp,
-    MakeRange,
-]
-
-NKI_ORIGINAL_OPS = {
-    ProgramId: nki_builder.program_id,
-    Allocate: nki_builder.ndarray,
-    Load: nki_builder.masked_load,
-    Store: nki_builder.masked_store,
-    Dot: nki_builder.matmul,
-    UnaryOp: nki_builder._unary_op,
-    MakeRange: nki_builder.arange,
-}
-
-NKI_OP_ATTR_NAMES = {
-    ProgramId: "program_id",
-    Allocate: "ndarray",
-    Load: "masked_load",
-    Store: "masked_store",
-    Dot: "matmul",
-    UnaryOp: "_unary_op",
-    MakeRange: "arange",
-}
+    NKI_OP_ATTR_NAMES = {
+        ProgramId: "program_id",
+        Allocate: "ndarray",
+        Load: "masked_load",
+        Store: "masked_store",
+        Dot: "matmul",
+        UnaryOp: "_unary_op",
+        MakeRange: "arange",
+    }
 
-NKI_ADAPTERS: dict[type[Op], Callable[..., AdapterResult]] = {
-    ProgramId: _program_id_adapter,
-    Allocate: _nki_allocate_adapter,
-    Load: _nki_load_adapter,
-    Store: _nki_store_adapter,
-    Dot: _nki_dot_adapter,
-}
+    NKI_ADAPTERS = {
+        ProgramId: _program_id_adapter,
+        Allocate: _nki_allocate_adapter,
+        Load: _nki_load_adapter,
+        Store: _nki_store_adapter,
+        Dot: _nki_dot_adapter,
+    }
 
-for op_type in NKI_OP_LIST:
-    NKI_ADAPTERS.setdefault(op_type, passthrough_adapter)
+    for op_type in NKI_OP_LIST:
+        NKI_ADAPTERS.setdefault(op_type, passthrough_adapter)
 
 
 OPERATION_REGISTRY: dict[str, dict[str, Any]] = {

From 0f6664ae8728be37c4222399c9c7d4207b391ba2 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Thu, 13 Nov 2025 10:49:18 -0500
Subject: [PATCH 83/90] guard nki test import if  NKI not installed

---
 tests/nki/test_nki.py | 9 ++++++++-
 1 file changed, 8 insertions(+), 1 deletion(-)

diff --git a/tests/nki/test_nki.py b/tests/nki/test_nki.py
index 01dc98cc..2ce3985f 100644
--- a/tests/nki/test_nki.py
+++ b/tests/nki/test_nki.py
@@ -3,9 +3,16 @@
 Test script to verify NDArray slicing functionality after fixes
 """
 import numpy as np
-from triton_viz.core.nki import NDArray
 import pytest
 
+try:
+    from triton_viz.core.nki import NDArray
+except ModuleNotFoundError:
+    pytest.skip(
+        "NeuronX dependencies are missing. Install triton-viz[nki] to run these tests.",
+        allow_module_level=True,
+    )
+
 pytestmark = pytest.mark.nki  # only run at "pytest -m nki"
 
 

From b45cd5e2f9b127af37a84a9de15b8a5799bcbe1e Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Thu, 13 Nov 2025 10:58:07 -0500
Subject: [PATCH 84/90] update import

---
 triton_viz/clients/sanitizer/sanitizer.py | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/triton_viz/clients/sanitizer/sanitizer.py b/triton_viz/clients/sanitizer/sanitizer.py
index 8ea4dee7..b00a04d6 100644
--- a/triton_viz/clients/sanitizer/sanitizer.py
+++ b/triton_viz/clients/sanitizer/sanitizer.py
@@ -1361,8 +1361,9 @@ def concretize(self_or_cls, obj=None):
                     obj.lhs.concretize(), obj.rhs.concretize(), obj.binary_numpy_op
                 )
         elif obj.op == "load":
-            from ...core.patch import original_ops
+            from ...core.patch import OPERATION_REGISTRY
 
+            original_ops = OPERATION_REGISTRY["triton"]["original_ops"]
             ptr_concrete = obj.ptr.concretize()
             # concretize mask
             # create an all-True mask if mask is None
@@ -1390,9 +1391,10 @@ def concretize(self_or_cls, obj=None):
         else:
             # Special handling for cast_impl and bitcast operations
             if obj.op in ("cast_impl", "bitcast"):
-                from ...core.patch import original_ops
+                from ...core.patch import OPERATION_REGISTRY
                 from ...core.data import CastImpl, Bitcast
 
+                original_ops = OPERATION_REGISTRY["triton"]["original_ops"]
                 src_concrete = obj.src.concretize()
                 # dst_type is stored as a SymbolicExpr const node, need to extract the value
                 dst_type_value = (

From 6c01d6299714f8d76ade161d0187c2d90756a238 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Thu, 13 Nov 2025 11:13:51 -0500
Subject: [PATCH 85/90] [CI] fix aws neuron ci

---
 .github/workflows/python-app.yml | 14 ++------------
 pyproject.toml                   |  2 --
 2 files changed, 2 insertions(+), 14 deletions(-)

diff --git a/.github/workflows/python-app.yml b/.github/workflows/python-app.yml
index d77ec62f..6e6781e3 100644
--- a/.github/workflows/python-app.yml
+++ b/.github/workflows/python-app.yml
@@ -80,19 +80,9 @@ jobs:
       with:
         python-version: '3.10'
 
-    - name: Configure AWS Neuron apt repository
-      run: |
-        sudo tee /etc/apt/sources.list.d/neuron.list <<'EOF'
-        deb https://apt.repos.neuron.amazonaws.com $(lsb_release -cs) main
-        EOF
-        wget -qO - https://apt.repos.neuron.amazonaws.com/GPG-PUB-KEY-AWS-NEURON.PUB | sudo apt-key add -
-        sudo apt-get update
-        sudo apt-get install -y aws-neuronx-runtime-lib aws-neuronx-collectives aws-neuronx-tools
-
     - name: Install AWS Neuron Python packages
       run: |
-        python -m pip install --upgrade pip
-        python -m pip install --upgrade neuronx-cc==2.* libneuronxla==2.* torch-neuronx==2.* torchvision --extra-index-url https://pip.repos.neuron.amazonaws.com
+        python -m pip install neuronx-cc==2.* --extra-index-url https://pip.repos.neuron.amazonaws.com
 
     - name: Install Triton
       run: |
@@ -106,4 +96,4 @@ jobs:
     - name: Run NKI pytest suite
       run: |
         cd triton_viz
-        python -m pytest tests/nki -m nki
+        python -m pytest tests/nki -m ""
diff --git a/pyproject.toml b/pyproject.toml
index 8680eb04..f3051005 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -50,6 +50,4 @@ triton-profiler = "triton_viz.wrapper:apply_profiler"
 [project.optional-dependencies]
 nki = [
   "neuronx-cc==2.*",
-  "libneuronxla==2.*",
-  "torch-neuronx==2.*",
 ]

From fa814b5e18f77d47d069e4b83fc5a7996f527373 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Thu, 13 Nov 2025 14:56:20 -0500
Subject: [PATCH 86/90] python>=3.10 needed (also 3.9 is EOL)

---
 README.md      | 9 +++++----
 pyproject.toml | 5 +++--
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/README.md b/README.md
index 372b1355..b29bf707 100644
--- a/README.md
+++ b/README.md
@@ -44,8 +44,9 @@ The best part about this tool is that while it does focus on visualizing GPU ope
 
 ### Prerequisites
 
-- Python installed (preferably the latest available version).
+- Python installed (preferably the latest available version), minimum supported version is 3.10.
 - [Triton](https://github.com/openai/triton/blob/main/README.md) installed. Follow the installation instructions in the linked repository.
+- Note: the below commands must be run in order.
 
 Upon successfully installing Triton, install Torch using the following command:
 
@@ -76,7 +77,7 @@ You're all set!
 If you want to exercise the Neuron Kernel Interface (NKI) interpreter or run the NKI-specific tests:
 
 1. Follow the [AWS Neuron Torch-NeuronX Ubuntu 22.04 setup guide](https://awsdocs-neuron.readthedocs-hosted.com/en/latest/setup/neuron-setup/pytorch/neuronx/ubuntu/torch-neuronx-ubuntu22.html#setup-torch-neuronx-ubuntu22) to add the Neuron APT repository and install the required system packages (for example `aws-neuronx-tools`, `aws-neuronx-runtime-lib`, `aws-neuronx-collectives`, and their dependencies).
-2. Install Triton-Viz with the optional NKI extras so the Neuron Python packages (`neuronx-cc`, `libneuronxla`, `torch-neuronx`) are available:
+2. Instead of running `pip install -e .` in the above section, install Triton-Viz with the optional NKI extras so the Neuron Python packages (`neuronx-cc`, `libneuronxla`, `torch-neuronx`) are available:
 
    ```sh
    pip install -e .[nki]
@@ -85,8 +86,8 @@ If you want to exercise the Neuron Kernel Interface (NKI) interpreter or run the
 
 ### Testing
 * To run core Triton-viz tests, run `pytest tests/`.
-* (if NKI installed) To run NKI-specific tests, run `pytest -m nki tests/`.
-* To run all tests (Triton + NKI), run `pytest -m "" tests/`.
+* (if NKI installed) To run NKI-specific tests, run `pytest tests/ -m nki`.
+* To run all tests (Triton + NKI), run `pytest tests/ -m ""`.
 
 ## Working with Examples
 
diff --git a/pyproject.toml b/pyproject.toml
index f3051005..dd0cf6e5 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -14,7 +14,7 @@ authors = [
 ]
 readme = "README.md"
 license = {text = "MIT"}
-requires-python = ">=3.7"
+requires-python = ">=3.10"
 classifiers = [
   "Programming Language :: Python :: 3",
   "License :: OSI Approved :: MIT License",
@@ -49,5 +49,6 @@ triton-profiler = "triton_viz.wrapper:apply_profiler"
 
 [project.optional-dependencies]
 nki = [
-  "neuronx-cc==2.*",
+  "neuronx-cc @ https://pip.repos.neuron.amazonaws.com/neuronx-cc/neuronx_cc-2.21.33363.0%2B82129205-cp310-cp310-linux_x86_64.whl ; python_version == '3.10'",
+  "neuronx-cc @ https://pip.repos.neuron.amazonaws.com/neuronx-cc/neuronx_cc-2.21.33363.0%2B82129205-cp311-cp311-linux_x86_64.whl ; python_version == '3.11'"
 ]

From de93383104ea18a6e913f002e354f5d8670e05c4 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Thu, 13 Nov 2025 14:57:08 -0500
Subject: [PATCH 87/90] [CI] make AWS tests reuse torch + triton installations
 from build job

---
 .github/workflows/python-app.yml | 33 +++-----------------------------
 1 file changed, 3 insertions(+), 30 deletions(-)

diff --git a/.github/workflows/python-app.yml b/.github/workflows/python-app.yml
index 6e6781e3..cb620c29 100644
--- a/.github/workflows/python-app.yml
+++ b/.github/workflows/python-app.yml
@@ -61,39 +61,12 @@ jobs:
         cd triton_viz
         python -m pytest tests
 
-  nki-tests:
-    name: AWS Neuron / NKI
-    runs-on: ubuntu-22.04
-    needs: build
-    env:
-      TRITON_INTERPRET: "1"
-      TRITON_VERSION: "3.4.0"
-      PIP_EXTRA_INDEX_URL: https://pip.repos.neuron.amazonaws.com
-
-    steps:
-    - uses: actions/checkout@v3
-      with:
-        path: 'triton_viz'
-
-    - name: Set up Python 3.10
-      uses: actions/setup-python@v3
-      with:
-        python-version: '3.10'
-
-    - name: Install AWS Neuron Python packages
-      run: |
-        python -m pip install neuronx-cc==2.* --extra-index-url https://pip.repos.neuron.amazonaws.com
-
-    - name: Install Triton
-      run: |
-        python -m pip install triton==${{ env.TRITON_VERSION }}
-
     - name: Install Triton-Viz with NKI extras
       run: |
         cd triton_viz
-        python -m pip install -e .[nki]
+        pip install -e .[nki]
 
-    - name: Run NKI pytest suite
+    - name: Run full (Triton + NKI) pytest suite
       run: |
         cd triton_viz
-        python -m pytest tests/nki -m ""
+        python -m pytest tests -m ""

From 50a0177a0209eabfed00e01c9728fb47c2449ba6 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Thu, 13 Nov 2025 16:12:03 -0500
Subject: [PATCH 88/90] make nki set_grid_dim consistent with triton api

---
 triton_viz/core/nki.py   | 4 ++--
 triton_viz/core/patch.py | 5 +----
 2 files changed, 3 insertions(+), 6 deletions(-)

diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index 9d7d97b2..fe797914 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -179,7 +179,7 @@ def __init__(self, grid_dims=None):
         self.fn = None
         self.shared_hbm_arrays = {}
 
-    def set_grid_dim(self, grid_dims):
+    def set_grid_dim(self, *grid_dims):
         self.grid_dims = grid_dims
 
     def set_grid_idx(self, x, y, z):
@@ -477,7 +477,7 @@ def run(self, *args, **kwargs):
             raise ValueError(
                 f"Grid must be 1, 2, or 3 dimensions, got {len(grid_dims)}"
             )
-        nki_builder.set_grid_dim(grid_dims)
+        nki_builder.set_grid_dim(*grid_dims)
         nki_builder.shared_hbm_arrays = {}
         nki_builder.fn = self.fn
 
diff --git a/triton_viz/core/patch.py b/triton_viz/core/patch.py
index f171c6f6..25186ff7 100644
--- a/triton_viz/core/patch.py
+++ b/triton_viz/core/patch.py
@@ -840,10 +840,7 @@ def run_grid_loops(grid):
     assert len(grid) <= 3
     grid = grid + (1,) * (3 - len(grid))
 
-    if backend == "nki":
-        builder.set_grid_dim(grid)
-    else:  # triton
-        builder.set_grid_dim(*grid)
+    builder.set_grid_dim(*grid)
     client_manager.grid_callback(grid)
     run_grid_loops(grid)
     # Copy arguments back to propagate side-effects

From 76f8999112f40f31d0b93965ad58ba1b9d9143d2 Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Fri, 28 Nov 2025 12:48:40 -0500
Subject: [PATCH 89/90] extra nki kernels to viz

---
 examples/nki/rmsnorm.py |  22 +--
 examples/nki/rope.py    | 294 ++++++++++++++++++++++++++++++++++++++++
 triton_viz/core/nki.py  |  10 ++
 3 files changed, 316 insertions(+), 10 deletions(-)
 create mode 100644 examples/nki/rope.py

diff --git a/examples/nki/rmsnorm.py b/examples/nki/rmsnorm.py
index da87ec0c..e844a934 100644
--- a/examples/nki/rmsnorm.py
+++ b/examples/nki/rmsnorm.py
@@ -19,9 +19,10 @@ def nki_rmsnorm_kernel(a_tensor, g_tensor, result):
     assert D == g_tensor.shape[0]
 
     # Generate tensor indices to index input tensor
-    ix = nl.arange(2)[:, None]
+    B_TILE = 8
+    ix = nl.arange(B_TILE)[:, None]
     iw = nl.arange(1)[:, None]
-    iy = nl.arange(8)[None, :]
+    iy = nl.arange(D)[None, :]
 
     # Load RMSNorm weight once, reused by rows/tiles of a_tensor
     g_tile = nl.load(g_tensor.reshape((1, D))[iw, iy], mask=iy < D)
@@ -29,10 +30,10 @@ def nki_rmsnorm_kernel(a_tensor, g_tensor, result):
     # Process 2 rows at a time due to 2-partition tile size limitation
     # Since we're not reducing across the first dimension
     # Tiles can be processed independently
-    for i in nl.affine_range(math.ceil(B / 2)):
+    for i in nl.affine_range(math.ceil(B / B_TILE)):
         # Load input data from external memory to on-chip memory
-        mask = (i * 2 + ix < B) & (iy < D)
-        a_tile = nl.load(a_tensor[i * 2 + ix, iy], mask=mask)
+        mask = (i * B_TILE + ix < B) & (iy < D)
+        a_tile = nl.load(a_tensor[i * B_TILE + ix, iy], mask=mask)
 
         # Compute element-wise square of a_tensor
         in_square = nl.square(a_tile)
@@ -53,13 +54,13 @@ def nki_rmsnorm_kernel(a_tensor, g_tensor, result):
 
         # Broadcast weight along first axis to match tensor shape
         # B_active = min(B - i * 2, 2)
-        g_bcast = g_tile.broadcast_to((2, 8))
+        g_bcast = g_tile.broadcast_to((B_TILE, D))
 
         # Multiply with the RMSNorm weight
-        out_tile = nl.multiply(out_tile, g_bcast, mask=(i * 2 + ix < B))
+        out_tile = nl.multiply(out_tile, g_bcast, mask=(i * B_TILE + ix < B))
 
         # store the addition results back to external memory (out_tensor)
-        nl.store(result[i * 2 + ix, iy], value=out_tile, mask=mask)
+        nl.store(result[i * B_TILE + ix, iy], value=out_tile, mask=mask)
 
 
 # ref
@@ -77,8 +78,9 @@ def torch_rmsnorm_kernel(a_tensor, g_tensor):
 
 TRITON_VIZ = True
 kernel_grid = (1, 1, 1)
-a_tensor = torch.arange(15).float().view(3, 5)
-g_tensor = torch.arange(5).float()
+B, D = 32, 32
+a_tensor = torch.arange(B * D).float().view(B, D)
+g_tensor = torch.arange(D).float()
 result = torch.empty_like(a_tensor).numpy()
 kernel_args = (a_tensor.numpy(), g_tensor.numpy(), result)
 
diff --git a/examples/nki/rope.py b/examples/nki/rope.py
new file mode 100644
index 00000000..0ec596ea
--- /dev/null
+++ b/examples/nki/rope.py
@@ -0,0 +1,294 @@
+import os
+from typing import Tuple
+
+import neuronxcc.nki.language as nl
+import torch
+from neuronxcc import nki
+import triton_viz
+from triton_viz.clients import Tracer
+from triton_viz.core.trace import launches
+import numpy as np
+
+os.environ["NEURON_FRAMEWORK_DEBUG"] = "1"
+# Ideally remove disable-dge for production performance, but kept for debugging context
+os.environ["NEURON_CC_FLAGS"] = " --disable-dge "
+
+
+def generate_pos_embedding(
+    head_dim: int, position_ids: torch.Tensor, base: int = 10000
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """
+    Generate positional embeddings for rotary position encoding (Llama style).
+    """
+    # Core RoPE block
+    inv_freq = 1.0 / (base ** (torch.arange(0, head_dim, 2) / head_dim))
+
+    # Expand to [HeadDim/2, 1]
+    inv_freq_expanded = inv_freq[:, None].float()
+    # Expand to [1, SeqLen]
+    position_ids_expanded = position_ids[None, :].float()
+
+    # MatMul -> [HeadDim/2, SeqLen] -> Transpose -> [SeqLen, HeadDim/2]
+    freqs = (inv_freq_expanded.float() @ position_ids_expanded.float()).transpose(0, 1)
+
+    # Concatenate to match HeadDim -> [Batch, SeqLen, HeadDim]
+    emb = torch.cat((freqs, freqs), dim=-1)
+    cos = emb.cos()
+    sin = emb.sin()
+    return cos, sin
+
+
+def div_ceil(n: int, d: int) -> int:
+    return (n + d - 1) // d
+
+
+def _nki_apply_rotary_embedding_core(q_tile, k_tile, cos_tile, sin_tile, output_tile):
+    """
+    Core NKI implementation of rotary position embedding computation.
+
+    Parameters
+    ----------
+    q_tile : nl.Tensor
+        Query tensor tile
+    k_tile : nl.Tensor
+        Key tensor tile
+    cos_tile : nl.Tensor
+        Cosine embedding tile
+    sin_tile : nl.Tensor
+        Sine embedding tile
+    output_tile : nl.Tensor
+        Output buffer for results
+
+    Notes
+    -----
+    The function applies rotary position embedding to query and key tensors
+    using the provided cosine and sine embeddings.
+    """
+
+    assert q_tile.shape[-1] % 2 == 0, "Sequence length for q_tile must be even!"
+    assert k_tile.shape[-1] % 2 == 0, "Sequence length for k_tile must be even!"
+    assert (
+        q_tile.shape[-1] == k_tile.shape[-1]
+    ), "q_tile and k_tile must have the same sequence length"
+
+    seq_len = q_tile.shape[-1]
+
+    # Rotate Q
+    output_tile[0, :, :] = q_tile * cos_tile
+    output_tile[0, :, : seq_len // 2] = output_tile[0, :, : seq_len // 2] + (
+        -1 * q_tile[:, seq_len // 2 :] * sin_tile[:, : seq_len // 2]
+    )
+    output_tile[0, :, seq_len // 2 :] = output_tile[0, :, seq_len // 2 :] + (
+        q_tile[:, : seq_len // 2] * sin_tile[:, seq_len // 2 :]
+    )
+
+    # Rotate K
+    output_tile[1, :, :] = k_tile * cos_tile
+    output_tile[1, :, : seq_len // 2] = output_tile[1, :, : seq_len // 2] + (
+        -1 * k_tile[:, seq_len // 2 :] * sin_tile[:, : seq_len // 2]
+    )
+    output_tile[1, :, seq_len // 2 :] = output_tile[1, :, seq_len // 2 :] + (
+        k_tile[:, : seq_len // 2] * sin_tile[:, seq_len // 2 :]
+    )
+
+
+def nki_rope_kernel(q, k, cos, sin, output_q, output_k):
+    """
+    NKI implementation of rotary position embedding.
+
+    Parameters
+    ----------
+    q : torch.Tensor
+        Query tensor of shape [batch_size, num_heads, seq_len, head_dim]
+    k : torch.Tensor
+        Key tensor of shape [batch_size, num_heads, seq_len, head_dim]
+    cos : torch.Tensor
+        Cosine embeddings
+    sin : torch.Tensor
+        Sine embeddings
+
+    Returns
+    -------
+    nl.Tensor
+        Output tensor containing transformed query and key tensors
+
+    Raises
+    ------
+    AssertionError
+        If input tensor shapes don't match or head dimension > 128
+    """
+    assert (
+        q.shape == k.shape
+    ), f"Shape of Q Tensor: {q.shape} doesn't match shape of K Tensor: {k.shape}"
+    assert (
+        cos.shape == sin.shape
+    ), f"Shape of cos Tensor: {cos.shape} doesn't match shape of sin Tensor: {sin.shape}"
+    PMAX = 128
+    assert (
+        q.shape[-1] <= PMAX
+    ), f"Shape of head dim (last dim) is more than PMAX: {q.shape}"
+
+    head_id = nl.program_id(axis=0)
+    seq_len = q.shape[1]
+    num_seq_batches = div_ceil(seq_len, nl.tile_size.pmax)
+    # output = nl.ndarray([2] + list(q.shape), dtype=q.dtype, buffer=nl.shared_hbm)
+    i_p, i_f = nl.mgrid[0:PMAX, 0 : q.shape[-1]]
+    # for seq_batch_id in nl.affine_range(0, num_seq_batches): # TODO
+    for seq_batch_id in nl.affine_range(num_seq_batches):
+        # q_hbm_tile = q[batch_id, head_id]
+        # k_hbm_tile = k[batch_id, head_id]
+        # cos_hbm_tile = cos[batch_id]
+        # sin_hbm_tile = sin[batch_id]
+
+        q_tile = nl.load(
+            # q_hbm_tile[seq_batch_id * nl.tile_size.pmax + i_p, i_f],
+            q[head_id, seq_batch_id * nl.tile_size.pmax + i_p, i_f],
+            mask=(seq_batch_id * nl.tile_size.pmax + i_p < seq_len),
+        )
+        k_tile = nl.load(
+            # k_hbm_tile[seq_batch_id * nl.tile_size.pmax + i_p, i_f],
+            k[head_id, seq_batch_id * nl.tile_size.pmax + i_p, i_f],
+            mask=(seq_batch_id * nl.tile_size.pmax + i_p < seq_len),
+        )
+        output_tile = nl.ndarray(
+            # [2] + [nl.par_dim(k_tile.shape[0]), k_tile.shape[1]], # TODO
+            [2] + [k_tile.shape[0], k_tile.shape[1]],
+            dtype=k_tile.dtype,
+            buffer=nl.sbuf,
+        )
+        cos_tile = nl.load(
+            # cos_hbm_tile[seq_batch_id * nl.tile_size.pmax + i_p, i_f],
+            cos[seq_batch_id * nl.tile_size.pmax + i_p, i_f],
+            mask=(seq_batch_id * nl.tile_size.pmax + i_p < seq_len),
+        )
+        sin_tile = nl.load(
+            # sin_hbm_tile[seq_batch_id * nl.tile_size.pmax + i_p, i_f],
+            sin[seq_batch_id * nl.tile_size.pmax + i_p, i_f],
+            mask=(seq_batch_id * nl.tile_size.pmax + i_p < seq_len),
+        )
+
+        _nki_apply_rotary_embedding_core(
+            q_tile, k_tile, cos_tile, sin_tile, output_tile
+        )
+
+        # output_q_hbm_tile = output[0, batch_id, head_id]
+        # output_k_hbm_tile = output[1, batch_id, head_id]
+
+        nl.store(
+            # output_q_hbm_tile[seq_batch_id * nl.tile_size.pmax + i_p, i_f],
+            output_q[head_id, seq_batch_id * nl.tile_size.pmax + i_p, i_f],
+            output_tile[0, :, :],
+            mask=(seq_batch_id * nl.tile_size.pmax + i_p < seq_len),
+        )
+        nl.store(
+            # output_k_hbm_tile[seq_batch_id * nl.tile_size.pmax + i_p, i_f],
+            output_k[head_id, seq_batch_id * nl.tile_size.pmax + i_p, i_f],
+            output_tile[1, :, :],
+            mask=(seq_batch_id * nl.tile_size.pmax + i_p < seq_len),
+        )
+
+    # return output
+
+
+# Torch reference implementation
+def torch_rope_kernel(q, k, cos, sin):
+    """Simple torch reference for rotary embedding"""
+
+    def rotate_half(x):
+        x1 = x[..., : x.shape[-1] // 2]
+        x2 = x[..., x.shape[-1] // 2 :]
+        return torch.cat((-x2, x1), dim=-1)
+
+    #
+    # FIX: Unsqueeze cos/sin to shape [B, 1, S, D] to broadcast over Heads [B, H, S, D]
+    cos = cos.unsqueeze(0)
+    sin = sin.unsqueeze(0)
+
+    q_rot = (q * cos) + (rotate_half(q) * sin)
+    k_rot = (k * cos) + (rotate_half(k) * sin)
+    return q_rot, k_rot
+
+
+# TRITON_VIZ test section
+TRITON_VIZ = True  # Set to True to see the trace
+# B, H, S, D = 2, 2, 4, 8
+H, S, D = 2, 4, 8
+# kernel_grid = (B, H)
+kernel_grid = (H,)
+
+if TRITON_VIZ:
+    # Setup test data
+    # q = torch.randn(B, H, S, D)
+    # k = torch.randn(B, H, S, D)
+    q = torch.randn(H, S, D)
+    k = torch.randn(H, S, D)
+    position_ids = torch.arange(S)
+    cos, sin = generate_pos_embedding(D, position_ids)
+
+    # output = torch.empty((2, B, H, S, D))
+    output_q = torch.empty((H, S, D))
+    output_k = torch.empty((H, S, D))
+    # Cast inputs to numpy for NKI tracing
+    kernel_args = (
+        q.numpy(),
+        k.numpy(),
+        cos.numpy(),
+        sin.numpy(),
+        output_q.numpy(),
+        output_k.numpy(),
+    )
+
+    print("Executing rotary embedding kernel with NKI interpreter...")
+    traced_kernel = triton_viz.trace(clients=Tracer(), backend="nki")(nki_rope_kernel)
+    kernel_instance = traced_kernel[kernel_grid]
+    kernel_instance(*kernel_args)
+
+    print(f"Number of launches: {len(launches)}")
+    if launches:
+        launch = launches[-1]
+        print(f"Number of records: {len(launch.records)}")
+        # Optional: Print record details
+
+    try:
+        triton_viz.launch(share=False)
+    except Exception as e:
+        print(f"Visualization error: {e}")
+
+else:
+    # Regular execution for comparison
+    # Ensure inputs are float32
+    q = torch.randn(H, S, D, dtype=torch.float32)
+    k = torch.randn(H, S, D, dtype=torch.float32)
+    position_ids = torch.arange(S)
+    cos, sin = generate_pos_embedding(D, position_ids)
+
+    print("Executing NKI JIT-ed rotary embedding kernel...")
+    compiled_kernel = nki.jit(nki_rope_kernel, kernel_return=False)
+    output_q_np = np.zeros((H, S, D), dtype=np.float32)
+    output_k_np = np.zeros((H, S, D), dtype=np.float32)
+
+    nki.simulate_kernel(
+        compiled_kernel[kernel_grid],
+        q.numpy(),
+        k.numpy(),
+        cos.numpy(),
+        sin.numpy(),
+        output_q_np,
+        output_k_np,
+    )
+
+    # Compare with torch reference
+    expected_q, expected_k = torch_rope_kernel(q, k, cos, sin)
+    actual_q = torch.from_numpy(output_q_np)
+    actual_k = torch.from_numpy(output_k_np)
+
+    max_diff_q = torch.max(torch.abs(expected_q - actual_q))
+    max_diff_k = torch.max(torch.abs(expected_k - actual_k))
+
+    print(f"Q max diff: {max_diff_q}")
+    print(f"K max diff: {max_diff_k}")
+
+    # Relaxed tolerance slightly for float32 accumulation differences
+    assert torch.allclose(expected_q, actual_q, rtol=1e-3, atol=1e-3), "Q mismatch"
+    assert torch.allclose(expected_k, actual_k, rtol=1e-3, atol=1e-3), "K mismatch"
+    print("Results match!")
diff --git a/triton_viz/core/nki.py b/triton_viz/core/nki.py
index fe797914..4517a251 100644
--- a/triton_viz/core/nki.py
+++ b/triton_viz/core/nki.py
@@ -85,6 +85,16 @@ def __getitem__(self, keys):
         # Create a new NDArray with the sliced data
         return NDArray(value=sliced_value, name=f"{self.name}_slice")
 
+    def __setitem__(self, keys, value):
+        if not isinstance(keys, tuple):
+            keys = (keys,)
+
+        # Apply the slicing to the underlying numpy array
+        new_keys = [k.data if isinstance(k, NDArray) else k for k in keys]
+        self.data[tuple(new_keys)] = value.data
+
+        return self
+
     def _binary_op(self, other, op_func, op_name, op_symbol):
         if isinstance(other, NDArray):
             return NDArray(

From 3dc8c020692d997648e4f1c5327757533d57db1e Mon Sep 17 00:00:00 2001
From: Thai-Hoa Nguyen <tonanhhoaduc@gmail.com>
Date: Fri, 5 Dec 2025 12:35:20 -0500
Subject: [PATCH 90/90] typo

---
 examples/nki/softmax.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/examples/nki/softmax.py b/examples/nki/softmax.py
index bd44b5e9..99942c0c 100644
--- a/examples/nki/softmax.py
+++ b/examples/nki/softmax.py
@@ -68,6 +68,6 @@ def softmax_kernel(in_tensor, out_tensor):
     compiled_kernel = nki.jit(softmax_kernel, kernel_return=False)
     nki.simulate_kernel(compiled_kernel[kernel_grid], *kernel_args)
 
-y_expected = np.exp(x_small) / np.exp(x_small).sum(-1, keepdim=True)
+y_expected = np.exp(x_small) / np.exp(x_small).sum(-1, keepdims=True)
 print(np.max(np.abs(y_expected - y_small)))
 assert np.allclose(y_expected, y_small)