[FIX] misc fixes for triton-puzzles (#252)

Author: latentCall145

.gitignore

@@ -54,7 +54,6 @@ coverage.xml
 .hypothesis/
 .pytest_cache/
 cover/
-triton_viz/static/
 
 # Translations
 *.mo
@@ -170,4 +169,3 @@ uv.lock
 triton_viz/version.py
 .subagents/
 subagent*.txt
-triton_viz/templates/index.html

README.md

@@ -67,7 +67,9 @@ If you want to run tests, run `uv sync --extra test` instead of `uv sync`. Other
 
 ### Frontend Build
 
-If you want to run the visualizer, build the TS sources:
+The PyPI package ships with prebuilt frontend assets in `triton_viz/static`, so
+you do not need npm to run the visualizer. If you want to modify the frontend,
+rebuild the TS sources:
 
 ```sh
 npm install

frontend/assets/visualizer.css

@@ -708,6 +708,21 @@ body {
     color: var(--text-secondary);
 }
 
+/* webgl fallback banner */
+.webgl-warning {
+    position: absolute;
+    inset: 24px;
+    display: flex;
+    align-items: center;
+    justify-content: center;
+    text-align: center;
+    font-size: 14px;
+    color: #f8fafc;
+    border-color: var(--warning);
+    background: rgba(15, 23, 42, 0.88);
+    box-shadow: var(--shadow-soft);
+}
+
 .is-draggable {
     cursor: grab;
     user-select: none;

frontend/components/tensor_view.ts

@@ -17,6 +17,8 @@ import {
     GAP,
     COLOR_HOVER,
     updateTensorHighlights,
+    canUseWebgl,
+    renderWebglWarning,
 } from '../utils/three_utils.js';
 import { createHistogramOverlay } from './histogram.js';
 import { enableDrag } from '../utils/ui_helpers.js';
@@ -852,14 +854,26 @@ export function createTensorVisualization(
         const stage = document.createElement('div');
         stage.className = 'viz-stage';
         containerElement.appendChild(stage);
+        if (!canUseWebgl()) {
+            // show a visible message when WebGL is disabled and skip initialization.
+            return renderWebglWarning(containerElement);
+        }
         const sideMenu = createSideMenu(containerElement);
         const histogramUI = createHistogramOverlay(containerElement, {
             title: `${type} Value Distribution`,
             apiBase: API_BASE,
             sources: configs.map(c => ({ value: c.name.toUpperCase(), label: `${c.name} Tensor` })),
             buildRequestBody: (s, b) => ({ uuid: op.uuid, source: s, bins: b }),
         });
-        const { scene, camera, renderer } = setupScene(stage, 0x000000);
+        let scene: ThreeScene;
+        let camera: ThreeCamera;
+        let renderer: ThreeRenderer;
+        try {
+            ({ scene, camera, renderer } = setupScene(stage, 0x000000));
+        } catch (err) {
+            // webgl can still fail even after a feature test.
+            return renderWebglWarning(containerElement);
+        }
         const disposer = createDisposer();
         const { cubeGeometry, edgesGeometry, lineMaterial } = setupGeometries();
         const tensors = new Map<string, TensorGroup>();

frontend/core/api.ts

@@ -9,7 +9,11 @@ export type RequestOptions = {
 const getDefaultBase = (): string => {
     if (typeof globalThis === 'undefined') return '';
     const globalBase = (globalThis as typeof globalThis & { __TRITON_VIZ_API__?: string }).__TRITON_VIZ_API__;
-    return globalBase || '';
+    if (globalBase) return globalBase;
+    if (typeof window === 'undefined' || !window.location) return '';
+    // derive a stable base from the current page location for proxy paths
+    const baseUrl = new URL('.', window.location.href);
+    return baseUrl.href.replace(/\/$/, '');
 };
 
 const buildUrl = (path?: string | null, baseOverride?: string | null): string => {

frontend/utils/three_utils.ts

@@ -22,6 +22,29 @@ export const COLOR_EDGE = new THREE.Color(0.5, 0.5, 0.5);
 
 const COLOR_SLICE = new THREE.Color(0.0, 0.7, 1.0);
 
+// quick feature test for WebGL availability in the current browser.
+export function canUseWebgl(): boolean {
+    if (typeof document === 'undefined') return false;
+    const canvas = document.createElement('canvas');
+    const gl = canvas.getContext('webgl2') || canvas.getContext('webgl') || canvas.getContext('experimental-webgl');
+    return !!gl;
+}
+
+// render a visible warning when WebGL is missing or disabled.
+export function renderWebglWarning(container: HTMLElement): () => void {
+    const existing = container.querySelector('.webgl-warning');
+    if (existing) {
+        return () => existing.remove();
+    }
+    const warning = document.createElement('div');
+    warning.className = 'webgl-warning info-card';
+    warning.setAttribute('role', 'alert');
+    warning.setAttribute('aria-live', 'assertive');
+    warning.textContent = 'WebGL is required for the 3D visualizer. Enable WebGL in your browser settings and reload this page.';
+    container.appendChild(warning);
+    return () => warning.remove();
+}
+
 export function setupScene(container: HTMLElement, backgroundColor = 0x000000): {
     scene: ThreeScene;
     camera: ThreeCamera;

pyproject.toml

@@ -34,7 +34,7 @@ dependencies = [
 ]
 
 [tool.setuptools.packages.find]
-include=["triton_viz"]
+include=["triton_viz*"]
 exclude=["tasks"]
 
 [tool.setuptools]

triton_viz/static/components/histogram.js

@@ -0,0 +1,181 @@
+import { postJson } from "../core/api.js";
+import { createDisposer } from "../utils/dispose.js";
+export function createHistogramOverlay(containerElement, options) {
+    const { title = "Value Distribution", sources = [], apiBase = "", buildRequestBody, defaultBins = 64, } = options;
+    if (!buildRequestBody) {
+        throw new Error("buildRequestBody is required for histogram overlay");
+    }
+    const button = document.createElement("button");
+    button.textContent = "Value Histogram";
+    const overlay = document.createElement("div");
+    Object.assign(overlay.style, {
+        position: "fixed",
+        top: "50px",
+        right: "50px",
+        width: "480px",
+        padding: "12px",
+        background: "rgba(0, 0, 0, 0.85)",
+        color: "#fff",
+        borderRadius: "8px",
+        border: "1px solid #555",
+        zIndex: 3000,
+        display: "none",
+    });
+    const header = document.createElement("div");
+    header.textContent = title;
+    header.style.fontSize = "16px";
+    header.style.marginBottom = "8px";
+    header.style.fontWeight = "bold";
+    overlay.appendChild(header);
+    const controls = document.createElement("div");
+    controls.style.display = "flex";
+    controls.style.gap = "12px";
+    controls.style.flexWrap = "wrap";
+    controls.style.alignItems = "flex-end";
+    const sourceGroup = document.createElement("div");
+    sourceGroup.style.display = "flex";
+    sourceGroup.style.flexDirection = "column";
+    sourceGroup.style.gap = "4px";
+    const sourceLabel = document.createElement("label");
+    sourceLabel.textContent = "Activation";
+    sourceLabel.style.fontSize = "12px";
+    sourceLabel.style.opacity = "0.8";
+    sourceGroup.appendChild(sourceLabel);
+    const select = document.createElement("select");
+    select.id = "histogram-source";
+    sources.forEach((src) => {
+        const opt = document.createElement("option");
+        opt.value = src.value;
+        opt.textContent = src.label;
+        select.appendChild(opt);
+    });
+    if (!select.value && select.options.length) {
+        select.selectedIndex = 0;
+    }
+    sourceLabel.htmlFor = select.id;
+    sourceGroup.appendChild(select);
+    controls.appendChild(sourceGroup);
+    const binsGroup = document.createElement("div");
+    binsGroup.style.display = "flex";
+    binsGroup.style.flexDirection = "column";
+    binsGroup.style.gap = "4px";
+    const binsLabel = document.createElement("label");
+    binsLabel.textContent = "Bins";
+    binsLabel.style.fontSize = "12px";
+    binsLabel.style.opacity = "0.8";
+    binsGroup.appendChild(binsLabel);
+    const binInput = document.createElement("input");
+    binInput.type = "number";
+    binInput.value = String(defaultBins);
+    binInput.min = "4";
+    binInput.max = "512";
+    binInput.step = "2";
+    binInput.style.width = "80px";
+    binInput.title = "Number of bins";
+    binInput.id = "histogram-bins";
+    binsLabel.htmlFor = binInput.id;
+    binsGroup.appendChild(binInput);
+    controls.appendChild(binsGroup);
+    overlay.appendChild(controls);
+    const info = document.createElement("div");
+    info.style.margin = "6px 0";
+    info.style.fontSize = "12px";
+    overlay.appendChild(info);
+    const canvas = document.createElement("canvas");
+    canvas.width = 440;
+    canvas.height = 240;
+    canvas.style.background = "#111";
+    canvas.style.border = "1px solid #444";
+    overlay.appendChild(canvas);
+    const status = document.createElement("div");
+    status.style.fontSize = "12px";
+    status.style.marginTop = "4px";
+    overlay.appendChild(status);
+    function show() {
+        overlay.style.display = "block";
+        updateHistogram();
+    }
+    function hide() {
+        overlay.style.display = "none";
+    }
+    const disposer = createDisposer();
+    disposer.listen(button, "click", show);
+    disposer.listen(select, "change", () => {
+        updateHistogram();
+    });
+    disposer.listen(binInput, "input", () => {
+        updateHistogram();
+    });
+    async function updateHistogram() {
+        status.textContent = "Loading histogram...";
+        info.textContent = "";
+        const bins = parseInt(binInput.value, 10) || defaultBins;
+        if (!select.value && select.options.length) {
+            select.selectedIndex = 0;
+        }
+        try {
+            const body = buildRequestBody(select.value, bins);
+            body.bins = bins;
+            body.max_samples = body.max_samples || 200000;
+            const data = await postJson("/api/histogram", body, { base: apiBase });
+            drawHistogram(canvas, data.counts, data.edges);
+            info.textContent = `Min: ${data.min.toFixed(6)} | Max: ${data.max.toFixed(6)} | Total values: ${data.n} | Sampled: ${data.sampled}`;
+            status.textContent = "";
+        }
+        catch (err) {
+            const message = err instanceof Error ? err.message : String(err);
+            status.textContent = `Histogram error: ${message}`;
+            const ctx = canvas.getContext("2d");
+            if (ctx)
+                ctx.clearRect(0, 0, canvas.width, canvas.height);
+        }
+    }
+    function drawHistogram(canvasEl, counts, edges) {
+        const ctx = canvasEl.getContext("2d");
+        if (!ctx)
+            return;
+        ctx.clearRect(0, 0, canvasEl.width, canvasEl.height);
+        if (!counts || !counts.length || !edges || edges.length < 2) {
+            ctx.fillStyle = "#888";
+            ctx.fillText("No data", 20, 30);
+            return;
+        }
+        const width = canvasEl.width - 40;
+        const height = canvasEl.height - 40;
+        const originX = 30;
+        const originY = canvasEl.height - 20;
+        const maxCount = Math.max(...counts);
+        const barWidth = width / counts.length;
+        ctx.strokeStyle = "#555";
+        ctx.beginPath();
+        ctx.moveTo(originX, originY);
+        ctx.lineTo(originX + width, originY);
+        ctx.stroke();
+        counts.forEach((count, idx) => {
+            const barHeight = maxCount ? (count / maxCount) * height : 0;
+            const x = originX + idx * barWidth;
+            const y = originY - barHeight;
+            ctx.fillStyle = "#ffa500";
+            ctx.fillRect(x, y, Math.max(1, barWidth - 2), barHeight);
+        });
+        ctx.fillStyle = "#ccc";
+        ctx.font = "10px monospace";
+        const firstEdge = edges[0];
+        const lastEdge = edges[edges.length - 1];
+        if (firstEdge === undefined || lastEdge === undefined)
+            return;
+        ctx.fillText(`${firstEdge.toFixed(4)}`, originX, originY + 12);
+        ctx.fillText(`${lastEdge.toFixed(4)}`, originX + width - 40, originY + 12);
+    }
+    (containerElement || document.body).appendChild(overlay);
+    return {
+        button,
+        overlay,
+        show,
+        hide,
+        destroy() {
+            disposer.dispose();
+            overlay.remove();
+        },
+    };
+}