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():