fix(inspect.histogram): handle rmsnorm / layernorm canvas nodes

Before: clicking 'Inspect weight histogram' on a rmsnorm or layernorm
node in BrickContextPanel returned the dispatcher's generic 'Invalid
params' chip because the handler raised ValueError when the brick's
kind was outside BLOCK_BUILDERS. Norm primitives live behind
_make_norm, not the block-builder registry, so any user with a norm
tail on the canvas (every tiny_aya / parallel-block preset) hit this.

Fix: branch in inspect_histogram —
  * BLOCK_BUILDERS[kind] → existing path
  * rmsnorm / layernorm → _make_norm(kind, hidden, eps)
  * residual → explicit 'no trainable weights' message instead of the
    generic 'unknown kind'

Bonus fix: dim_env is a dict in the wire schema, but the handler was
using getattr(dim_env, "H", 128) which silently always returned 128
since dicts don't have attributes. Switched to .get('H', 128).

Verified end-to-end: aya_norm (rmsnorm @ H=128) now returns 128
weight values with mean = 1.0 (γ init).

Tests: 3 new (rmsnorm, layernorm, residual rejection). Total
inspect_histogram coverage 6/6 green.

Author: apstenku123

cppmega_v4/jsonrpc/histogram_method.py

@@ -89,13 +89,34 @@ def inspect_histogram(params: HistogramParams,
             f"brick_id={params.brick_id!r} not in spec.graph.nodes"
         )
     kind = getattr(brick_node, "kind", "")
-    if kind not in BLOCK_BUILDERS:
+    # dim_env is a dict in the wire schema — use mapping access, not
+    # getattr (which would always return the default for a dict).
+    dim_env = params.spec.dim_env or {}
+    hidden = int(dim_env.get("H", 128)) if isinstance(dim_env, dict) \
+        else int(getattr(dim_env, "H", 128))
+    brick_params = dict(getattr(brick_node, "params", {}) or {})
+    if kind in BLOCK_BUILDERS:
+        module = BLOCK_BUILDERS[kind](hidden, brick_params)
+    elif kind in {"rmsnorm", "layernorm"}:
+        # Norm primitives aren't in BLOCK_BUILDERS — instantiate
+        # directly so 'Inspect weight histogram' works on canvas norm
+        # nodes (e.g. tiny_aya rmsnorm tail). eps defaults to 1e-5.
+        from cppmega_v4.models.unified_superblock_v4 import _make_norm
+        eps = float(brick_params.get("eps", 1e-5))
+        module = _make_norm(kind, hidden, eps)
+        if module is None:
+            raise ValueError(
+                f"brick_id={params.brick_id!r} norm kind={kind!r} "
+                f"resolved to a no-op (no weights to inspect)")
+    elif kind == "residual":
+        # Pure plumbing node — no parameters to histogram.
+        raise ValueError(
+            f"brick_id={params.brick_id!r} kind='residual' has no "
+            f"trainable weights (residual is a join, not a module)")
+    else:
         raise ValueError(
             f"brick_id={params.brick_id!r} has unknown kind={kind!r}"
         )
-    hidden = int(getattr(params.spec.dim_env, "H", 128))
-    brick_params = dict(getattr(brick_node, "params", {}) or {})
-    module = BLOCK_BUILDERS[kind](hidden, brick_params)
     flat_parts: list[mx.array] = []
     for _k, v in nn.utils.tree_flatten(module.parameters()):
         if hasattr(v, "shape"):

tests/v4/test_inspect_histogram.py

@@ -56,3 +56,91 @@ def test_v7_h08_bucket_clip_64():
     # Bins are monotonic.
     for a, b in zip(r.bins, r.bins[1:]):
         assert a < b
+
+
+def test_histogram_handles_rmsnorm_node():
+    """fhxg: Inspect weight histogram on a `rmsnorm` canvas node must
+    succeed — the kind is a norm primitive (not in BLOCK_BUILDERS) and
+    used to surface 'Invalid params' before the fallback was added.
+
+    RMSNorm has `dim` channels initialised to 1.0 (gamma), so the
+    histogram is a single peak at 1.0.
+    """
+    from cppmega_v4.jsonrpc.histogram_method import (
+        HistogramParams, inspect_histogram,
+    )
+    from cppmega_v4.jsonrpc.schema import VerifyParams
+    spec = VerifyParams.model_validate({
+        "graph": {
+            "nodes": [
+                {"id": "a", "kind": "attention", "params": {}},
+                {"id": "n", "kind": "rmsnorm", "params": {}},
+            ],
+            "edges": [{"src": "a", "dst": "n"}],
+        },
+        "dim_env": {"H": 128},
+        "loss": {"kind": "cross_entropy", "head_outputs": ["n"]},
+        "optim": {"kind": "adamw", "groups": [
+            {"matcher": "all", "lr": 1e-3, "weight_decay": 0.01,
+             "betas": [0.9, 0.95]}]},
+        "sharding": None,
+        "training": True,
+    })
+    r = inspect_histogram(HistogramParams(
+        spec=spec, brick_id="n", kind="weight", buckets=32))
+    assert r.n_values == 128
+    assert abs(r.mean - 1.0) < 1e-6
+
+
+def test_histogram_handles_layernorm_node():
+    from cppmega_v4.jsonrpc.histogram_method import (
+        HistogramParams, inspect_histogram,
+    )
+    from cppmega_v4.jsonrpc.schema import VerifyParams
+    spec = VerifyParams.model_validate({
+        "graph": {
+            "nodes": [
+                {"id": "a", "kind": "attention", "params": {}},
+                {"id": "n", "kind": "layernorm", "params": {}},
+            ],
+            "edges": [{"src": "a", "dst": "n"}],
+        },
+        "dim_env": {"H": 64},
+        "loss": {"kind": "cross_entropy", "head_outputs": ["n"]},
+        "optim": {"kind": "adamw", "groups": [
+            {"matcher": "all", "lr": 1e-3, "weight_decay": 0.01,
+             "betas": [0.9, 0.95]}]},
+        "sharding": None,
+        "training": True,
+    })
+    r = inspect_histogram(HistogramParams(
+        spec=spec, brick_id="n", kind="weight", buckets=16))
+    # LayerNorm has gamma + beta = 2 * H channels.
+    assert r.n_values == 2 * 64
+
+
+def test_histogram_residual_kind_rejected_cleanly():
+    import pytest
+    from cppmega_v4.jsonrpc.histogram_method import (
+        HistogramParams, inspect_histogram,
+    )
+    from cppmega_v4.jsonrpc.schema import VerifyParams
+    spec = VerifyParams.model_validate({
+        "graph": {
+            "nodes": [
+                {"id": "a", "kind": "attention", "params": {}},
+                {"id": "r", "kind": "residual", "params": {}},
+            ],
+            "edges": [{"src": "a", "dst": "r"}],
+        },
+        "dim_env": {"H": 64},
+        "loss": {"kind": "cross_entropy", "head_outputs": ["r"]},
+        "optim": {"kind": "adamw", "groups": [
+            {"matcher": "all", "lr": 1e-3, "weight_decay": 0.01,
+             "betas": [0.9, 0.95]}]},
+        "sharding": None,
+        "training": True,
+    })
+    with pytest.raises(ValueError, match="no trainable weights"):
+        inspect_histogram(HistogramParams(
+            spec=spec, brick_id="r", kind="weight", buckets=8))