analytical-viz: consolidate shape tables under Physical Layout tab

Move the per-stage shape tables (attention + FFN) off the Per-stage
latency tab and combine them with the weight + KV shape tables into a
single 'All tensor shapes (per PE)' expander at the bottom of the
Physical Layout tab. Five subsections in one place:
  - Attention weights (global + per-PE shape, per layer)
  - FFN weights (same schema)
  - KV cache (same schema)
  - Per-stage attention shapes (input / weight / output)
  - Per-stage FFN shapes (same)

Rationale: the shape catalog belongs next to the physical-layout view
that shows how the tensors are placed on hardware. The Per-stage
latency tab now focuses purely on time-per-stage. Weight/KV tables
stay on the Memory Breakdown tab too because the summary metrics
underneath still consume the same rows.

Guard test verifies the section string appears exactly once, is
positioned inside tab_layout (before tab_memory / tab_stages), and
that neither attn_stage_shape_rows nor ffn_stage_shape_rows is called
inside the tab_stages block anymore.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>

tests/analytical_visualization/app.py

@@ -1058,6 +1058,52 @@ with tab_layout:
             })
         st.dataframe(pd.DataFrame(rows), width='stretch', hide_index=True)
 
+    # Row G: All tensor shapes (per PE) — one collapsed expander with
+    # weight, KV, and per-stage shapes for anyone who wants them.
+    with st.expander("All tensor shapes (per PE)", expanded=False):
+        st.caption(
+            "Every tensor this deployment holds, per PE. Weights + KV "
+            "show global vs per-PE shape (one row per tensor, per layer). "
+            "Per-stage shapes show input / weight / output tensors each "
+            "stage of the forward pass operates on."
+        )
+        _lp = (model.layers + topo.pp - 1) // topo.pp
+
+        st.markdown(f"**Attention weights** (per layer, layers/stage = {_lp})")
+        _attn_rows_layout = attention_weight_rows(cfg)
+        _display_attn_layout = [
+            {k: v for k, v in r.items() if not k.startswith("_")}
+            for r in _attn_rows_layout
+        ]
+        st.dataframe(pd.DataFrame(_display_attn_layout),
+                     width='stretch', hide_index=True)
+
+        st.markdown(f"**FFN weights** (per layer, layers/stage = {_lp})")
+        _ffn_rows_layout = ffn_weight_rows(cfg)
+        _display_ffn_layout = [
+            {k: v for k, v in r.items() if not k.startswith("_")}
+            for r in _ffn_rows_layout
+        ]
+        st.dataframe(pd.DataFrame(_display_ffn_layout),
+                     width='stretch', hide_index=True)
+
+        st.markdown(f"**KV cache** (per layer, S_kv = {s_kv:,})")
+        _kv_rows_layout = kv_cache_rows(cfg)
+        _display_kv_layout = [
+            {k: v for k, v in r.items() if not k.startswith("_")}
+            for r in _kv_rows_layout
+        ]
+        st.dataframe(pd.DataFrame(_display_kv_layout),
+                     width='stretch', hide_index=True)
+
+        st.markdown("**Per-stage attention shapes**")
+        st.dataframe(pd.DataFrame(attn_stage_shape_rows(cfg)),
+                     width='stretch', hide_index=True)
+
+        st.markdown("**Per-stage FFN shapes**")
+        st.dataframe(pd.DataFrame(ffn_stage_shape_rows(cfg)),
+                     width='stretch', hide_index=True)
+
 
 # ── TAB 2: Memory breakdown ─────────────────────────────────────
 with tab_memory:
@@ -1227,22 +1273,6 @@ with tab_stages:
     st.pyplot(fig_bar, width='stretch')
     plt.close(fig_bar)
 
-    # ── Per-stage tensor shapes (per PE) ───────────────────────────
-    st.divider()
-    st.subheader("Per-stage tensor shapes (per PE)")
-    st.caption(
-        "Input / weight / output tensor dimensions each stage operates on, "
-        "*per PE* (already sharded by CP, TP, PP, EP as configured). Rows "
-        "match the latency table above; comm-only stages show payload in "
-        "the Input column."
-    )
-    st.markdown("**Attention stages**")
-    st.dataframe(pd.DataFrame(attn_stage_shape_rows(cfg)),
-                 width='stretch', hide_index=True)
-    st.markdown("**FFN stages**")
-    st.dataframe(pd.DataFrame(ffn_stage_shape_rows(cfg)),
-                 width='stretch', hide_index=True)
-
 
 # ── TAB 4: Save & compare ───────────────────────────────────────
 def _snapshot_current_config():

tests/analytical_visualization/test_stage_shapes.py

@@ -113,3 +113,42 @@ def test_ffn_gate_up_output_matches_ffn_per_pe():
     rows = ffn_stage_shape_rows(cfg)
     f2 = next(r for r in rows if r["Stage"] == "F2")
     assert str(ffn_pe) in f2["Output (per PE)"], (ffn_pe, f2)
+
+
+# ── App wiring: consolidated shapes section lives on layout tab ───
+
+
+def test_shapes_section_is_on_physical_layout_tab():
+    """The 'All tensor shapes (per PE)' expander must live inside
+    tab_layout (Physical layout), not tab_stages. Guards the move so
+    the section doesn't accidentally end up back on the latency tab.
+    """
+    from pathlib import Path
+    app_path = (
+        Path(__file__).parent / "app.py"
+    ).resolve()
+    src = app_path.read_text(encoding="utf-8")
+    # Section appears exactly once.
+    assert src.count('"All tensor shapes (per PE)"') == 1, src.count(
+        '"All tensor shapes (per PE)"'
+    )
+    # And it sits inside tab_layout, before tab_memory begins.
+    section_pos = src.index('"All tensor shapes (per PE)"')
+    tab_memory_pos = src.index("with tab_memory:")
+    tab_stages_pos = src.index("with tab_stages:")
+    assert section_pos < tab_memory_pos, (
+        "shapes section must be inside tab_layout (before tab_memory)"
+    )
+    assert section_pos < tab_stages_pos, (
+        "shapes section must be inside tab_layout (before tab_stages)"
+    )
+    # The old per-stage shape rendering under tab_stages must be gone —
+    # no attn_stage_shape_rows / ffn_stage_shape_rows call may sit
+    # between 'with tab_stages:' and the next 'with tab_' block.
+    stages_block = src[tab_stages_pos:src.index("# ── TAB", tab_stages_pos + 1)]
+    assert "attn_stage_shape_rows(cfg)" not in stages_block, (
+        "per-stage shape table must not remain on Per-stage latency tab"
+    )
+    assert "ffn_stage_shape_rows(cfg)" not in stages_block, (
+        "per-stage FFN shape table must not remain on Per-stage latency tab"
+    )