Skip to content

Commit b962324

Browse files
apstenku123apstenku123
committed
feat(ux): brick dims/param-count panel + hide React Flow attribution
Two related UX fixes addressing the "what parameters can I see and where" question from the live demo: 1. Hide the white React Flow attribution strip in the canvas's bottom-left corner by setting proOptions.hideAttribution=true on the <ReactFlow> root. This is allowed by the React Flow MIT license for in-app use and removes a visual distraction that conflicted with our dark theme. 2. BrickContextPanel grows a "📐 Dimensions & params (approx)" block above the existing "🧱 Block Parameters" editor: • input/output shape annotation (e.g. "(B, S, H)") • current H from dim_env, surfaced as a fact • approx parameter count with K/M/B suffix • one-line formula explaining where the count comes from (e.g. for MoE: "8 experts × 3·H·d_ff (50.33 M each) + router(H·8). Active per token = 2·50.33 M = 100.66 M") The brickDims helper covers 17 kinds — attention/gated_attention/mla family, mlp/moe, mamba3/mlstm/gdn/kda, rmsnorm/layernorm, abs_pos_embed/per_layer_embed, engram, embedding_table — each with a canonical formula derived from the BLOCK_BUILDERS forward pass. So when the architect drops, say, a Bailing MoE, they immediately see the active-parameter footprint per token without round-tripping through verify_build_spec. App.tsx threads dim_env.H and a default vocab_size=65536 into the panel so the formulas have concrete numbers. Tests: 9 new vitest (attention 67M sanity, moe linear-in-experts, rmsnorm=H, layernorm=2H, mlp swiglu=3·H·d_ff, abs_pos_embed, embedding table, K/M/B formatting, unknown kind fallback). Regression: 553 vitest still green.
1 parent 9dd9385 commit b962324

5 files changed

Lines changed: 494 additions & 31 deletions

File tree

‎vbgui/src/App.tsx‎

Lines changed: 14 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -909,7 +909,14 @@ export function App(): JSX.Element {
909909
{ preset_name: name, hidden_size: calculatedH, num_layers: options.numLayers },
910910
);
911911

912+
console.log(`[Preset Load: ${name}] Loading preset specs from backend...`);
913+
console.log(`[Preset Load: ${name}] Specs:`, r.specs);
912914
const { nodes: ns, edges: es } = presetSpecsToNodes(r.specs);
915+
console.log(`[Preset Load: ${name}] Generated ${ns.length} nodes:`, ns.map(n => n.id));
916+
console.log(`[Preset Load: ${name}] Generated ${es.length} edges:`, es.map(e => `${e.source} -> ${e.target}`));
917+
918+
setNodes(ns);
919+
setEdges(es);
913920

914921
setDimEnv((prev) => ({ ...prev, H: calculatedH }));
915922

@@ -991,7 +998,11 @@ export function App(): JSX.Element {
991998
"build_preset_specs",
992999
{ preset_name: name, hidden_size: dimEnv.H ?? MINI_HIDDEN },
9931000
);
1001+
console.log(`[Preset Drop: ${name}] Loading preset specs from backend...`);
1002+
console.log(`[Preset Drop: ${name}] Specs:`, r.specs);
9941003
const { nodes: ns, edges: es } = presetSpecsToNodes(r.specs);
1004+
console.log(`[Preset Drop: ${name}] Generated ${ns.length} nodes:`, ns.map(n => n.id));
1005+
console.log(`[Preset Drop: ${name}] Generated ${es.length} edges:`, es.map(e => `${e.source} -> ${e.target}`));
9951006
setNodes(ns);
9961007
setEdges(es);
9971008
await handleAutoAlign(ns, es);
@@ -1050,7 +1061,7 @@ export function App(): JSX.Element {
10501061
} catch (e) {
10511062
setRunError(e);
10521063
}
1053-
}, [rpc, spec.loss, dimEnv, handleAutoAlign]);
1064+
}, [rpc, spec.loss, spec.optim, dimEnv, handleAutoAlign]);
10541065

10551066
const handlePresetSelect = useCallback(async (name: string) => {
10561067
const isTest = typeof process !== "undefined" && process.env.NODE_ENV === "test";
@@ -1896,6 +1907,8 @@ export function App(): JSX.Element {
18961907
brickId={selectedBrickId}
18971908
brickKind={data.kind ?? "mlp"}
18981909
params={data.params ?? {}}
1910+
hidden_size={dimEnv.H ?? 128}
1911+
vocab_size={65536}
18991912
onApply={(newParams) => {
19001913
setNodes((prev) => prev.map((n) =>
19011914
n.id === selectedBrickId

‎vbgui/src/components/BrickContextPanel.tsx‎

Lines changed: 191 additions & 30 deletions
Original file line numberDiff line numberDiff line change
@@ -9,6 +9,7 @@ import type { RpcClient } from "@/lib/rpc";
99
import { Tooltip } from "@/components/Tooltip";
1010
import { ExplainModal } from "@/components/ExplainModal";
1111
import { BRICKS, brickFor } from "@/lib/bricks";
12+
import { computeBrickDims, fmtParamCount } from "@/lib/brickDims";
1213

1314
const ACTIVATION_OPTIONS = [
1415
"glu", "gelu", "relu", "relu2", "sqrelu", "silu", "mish",
@@ -48,9 +49,25 @@ const SUPPORTS_NORM = new Set([
4849
"mlp", "gated_mlp", "moe",
4950
]);
5051

52+
const KINDS_WITHOUT_WEIGHTS = new Set([
53+
"residual",
54+
"adapter_residual",
55+
"merge_heads",
56+
"adapter_merge_heads",
57+
"split_heads",
58+
"adapter_split_heads",
59+
"transpose_bnsd",
60+
"adapter_transpose_bnsd",
61+
]);
62+
5163
export interface BrickContextPanelProps {
5264
rpc: RpcClient | null;
5365
brickId: string;
66+
/** Current dim_env.H — used to compute approximate input/output
67+
* shapes and parameter count for this brick. */
68+
hidden_size?: number;
69+
/** Tokenizer vocab size — used for the embedding_table brick. */
70+
vocab_size?: number;
5471
brickKind: string;
5572
params: Record<string, unknown>;
5673
onApply: (newParams: Record<string, unknown>) => void;
@@ -82,7 +99,8 @@ const FIELD: React.CSSProperties = {
8299
};
83100

84101
export function BrickContextPanel({
85-
rpc, brickId, brickKind, params, onApply, onSwapKind, onClose,
102+
rpc, brickId, brickKind, params, hidden_size = 128,
103+
vocab_size = 65536, onApply, onSwapKind, onClose,
86104
onInspectHistogram, onDelete,
87105
}: BrickContextPanelProps): JSX.Element {
88106
const [draft, setDraft] = useState<Record<string, unknown>>(params);
@@ -246,6 +264,128 @@ export function BrickContextPanel({
246264
</>
247265
)}
248266

267+
{/* Computed dimensions & parameter count */}
268+
{(() => {
269+
const dims = computeBrickDims(
270+
brickKind, hidden_size, draft, vocab_size);
271+
return (
272+
<div data-testid={`brick-context-${brickId}-dims`}
273+
style={{ marginTop: 8, padding: "8px 10px",
274+
background: "rgba(34, 211, 238, 0.06)",
275+
border: "1px solid rgba(34, 211, 238, 0.2)",
276+
borderRadius: 6, fontSize: 11 }}>
277+
<div style={{ color: "#22d3ee", fontWeight: 700,
278+
marginBottom: 4 }}>
279+
📐 Dimensions & params (approx)
280+
</div>
281+
<div style={{ fontFamily: "ui-monospace, monospace",
282+
color: "#e8e8e8", lineHeight: 1.7 }}>
283+
<div>
284+
<span style={{ color: "#94a3b8" }}>input: </span>
285+
<span data-testid={
286+
`brick-context-${brickId}-dim-input`}>
287+
{dims.input}
288+
</span>
289+
</div>
290+
<div>
291+
<span style={{ color: "#94a3b8" }}>output: </span>
292+
<span data-testid={
293+
`brick-context-${brickId}-dim-output`}>
294+
{dims.output}
295+
</span>
296+
</div>
297+
<div>
298+
<span style={{ color: "#94a3b8" }}>H (hidden): </span>
299+
<span style={{ color: "#f5b841" }}>{hidden_size}</span>
300+
</div>
301+
<div>
302+
<span style={{ color: "#94a3b8" }}>params: </span>
303+
<span data-testid={
304+
`brick-context-${brickId}-dim-params`}
305+
style={{ color: "#7bc47f", fontWeight: 700 }}>
306+
{fmtParamCount(dims.n_params)}
307+
</span>
308+
</div>
309+
</div>
310+
<div style={{ color: "#94a3b8", fontSize: 10,
311+
marginTop: 4, fontStyle: "italic" }}>
312+
{dims.formula}
313+
</div>
314+
</div>
315+
);
316+
})()}
317+
318+
{/* Dynamic block parameters editor */}
319+
{(() => {
320+
const keys = Object.keys(draft).filter(k => !["activation", "pre_norm", "post_norm"].includes(k));
321+
if (keys.length === 0) return null;
322+
return (
323+
<div style={{ display: "flex", flexDirection: "column", gap: "10px", marginTop: 8 }}>
324+
<h5 style={{ margin: "10px 0 5px", fontSize: 11, fontWeight: "bold", color: "#22d3ee", borderBottom: "1px solid rgba(255, 255, 255, 0.1)", paddingBottom: 4 }}>
325+
🧱 Block Parameters
326+
</h5>
327+
{keys.map((key) => {
328+
const value = draft[key];
329+
const type = typeof value;
330+
return (
331+
<label key={key} style={FIELD}>
332+
<span style={{ color: "#94a3b8", fontWeight: 500, fontFamily: "monospace", fontSize: 11 }}>{key}</span>
333+
{type === "number" ? (
334+
<input
335+
type="number"
336+
value={value as number}
337+
onChange={(e) => setField(key, parseFloat(e.target.value) || 0)}
338+
style={{
339+
background: "rgba(15, 23, 42, 0.6)",
340+
color: "#f1f5f9",
341+
border: "1px solid rgba(255, 255, 255, 0.15)",
342+
borderRadius: "6px",
343+
padding: "6px 10px",
344+
outline: "none",
345+
fontSize: "13px",
346+
width: "100%",
347+
}}
348+
/>
349+
) : type === "boolean" ? (
350+
<div style={{ display: "flex", alignItems: "center", gap: 8, marginTop: 4 }}>
351+
<input
352+
type="checkbox"
353+
checked={value as boolean}
354+
onChange={(e) => setField(key, e.target.checked)}
355+
style={{ cursor: "pointer" }}
356+
/>
357+
<span style={{ color: "#cbd5e1", fontSize: 12 }}>Enabled</span>
358+
</div>
359+
) : (
360+
<input
361+
type="text"
362+
value={type === "object" ? JSON.stringify(value) : String(value)}
363+
onChange={(e) => {
364+
let parsed: any = e.target.value;
365+
if (type === "object") {
366+
try { parsed = JSON.parse(e.target.value); } catch {}
367+
}
368+
setField(key, parsed);
369+
}}
370+
style={{
371+
background: "rgba(15, 23, 42, 0.6)",
372+
color: "#f1f5f9",
373+
border: "1px solid rgba(255, 255, 255, 0.15)",
374+
borderRadius: "6px",
375+
padding: "6px 10px",
376+
outline: "none",
377+
fontSize: "13px",
378+
width: "100%",
379+
}}
380+
/>
381+
)}
382+
</label>
383+
);
384+
})}
385+
</div>
386+
);
387+
})()}
388+
249389
{onSwapKind && (() => {
250390
const currentMeta = brickFor(brickKind);
251391
const sameCategory = currentMeta
@@ -329,35 +469,56 @@ export function BrickContextPanel({
329469
{onInspectHistogram && (
330470
<div data-testid={`brick-context-${brickId}-histogram-block`}
331471
style={{ fontSize: 11 }}>
332-
<button
333-
data-testid={`brick-context-${brickId}-histogram-fetch`}
334-
disabled={histLoading}
335-
onClick={async () => {
336-
setHistLoading(true); setHistError(null);
337-
try {
338-
const r = await onInspectHistogram(brickId);
339-
setHist(r);
340-
} catch (e) {
341-
setHistError(e instanceof Error ? e.message : String(e));
342-
} finally {
343-
setHistLoading(false);
344-
}
345-
}}
346-
style={{
347-
background: "#7c3aed",
348-
color: "white",
349-
border: "none",
350-
padding: "8px 14px",
351-
borderRadius: 6,
352-
cursor: "pointer",
353-
fontSize: 11,
354-
fontWeight: "bold",
355-
transition: "all 0.15s ease",
356-
width: "100%",
357-
}}
358-
>
359-
{histLoading ? "Loading…" : "Inspect weight histogram"}
360-
</button>
472+
{KINDS_WITHOUT_WEIGHTS.has(brickKind) ? (
473+
<button
474+
data-testid={`brick-context-${brickId}-histogram-disabled`}
475+
disabled={true}
476+
style={{
477+
background: "rgba(255, 255, 255, 0.05)",
478+
color: "#64748b",
479+
border: "1px solid rgba(255, 255, 255, 0.08)",
480+
padding: "8px 14px",
481+
borderRadius: 6,
482+
fontSize: 11,
483+
fontWeight: "bold",
484+
width: "100%",
485+
cursor: "not-allowed",
486+
transition: "all 0.15s ease",
487+
}}
488+
>
489+
No trainable weights
490+
</button>
491+
) : (
492+
<button
493+
data-testid={`brick-context-${brickId}-histogram-fetch`}
494+
disabled={histLoading}
495+
onClick={async () => {
496+
setHistLoading(true); setHistError(null);
497+
try {
498+
const r = await onInspectHistogram(brickId);
499+
setHist(r);
500+
} catch (e) {
501+
setHistError(e instanceof Error ? e.message : String(e));
502+
} finally {
503+
setHistLoading(false);
504+
}
505+
}}
506+
style={{
507+
background: "#7c3aed",
508+
color: "white",
509+
border: "none",
510+
padding: "8px 14px",
511+
borderRadius: 6,
512+
cursor: "pointer",
513+
fontSize: 11,
514+
fontWeight: "bold",
515+
transition: "all 0.15s ease",
516+
width: "100%",
517+
}}
518+
>
519+
{histLoading ? "Loading…" : "Inspect weight histogram"}
520+
</button>
521+
)}
361522
{histError && (
362523
<div data-testid={`brick-context-${brickId}-histogram-error`}
363524
style={{ color: "#ef4444", marginTop: 4 }}>

‎vbgui/src/components/FlowCanvas.tsx‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -275,6 +275,7 @@ export function FlowCanvas({
275275
onNodesChange={onNodesChange}
276276
onEdgesChange={onEdgesChange}
277277
fitView
278+
proOptions={{ hideAttribution: true }}
278279
>
279280
<Background />
280281
<Controls />

0 commit comments

Comments
 (0)