Hidden state capture layer hook

Skill kjuhwa/skills-hub/skills/llm-agents/hidden-state-capture-layer-hook

Self-correcting knowledge corpus for Claude Code — 9 stable shape clusters, bias-correction pipeline baked into contribution flow. 47 papers, 45 techniques, 1.1k skills.

Install

npx -y skills add kjuhwa/skills-hub --skill hidden-state-capture-layer-hook

Assembled from the repository path, not quoted from the project. Check it against their README if it does not work.

One thing to look at

0 stars0 stars. Stars are a popularity signal and not a quality one, but at this level it is likely that nobody has read this closely except its author, and you would be relying on your own review.

What its author says it does

Copied from the file, not written here

Capture hidden states from specified transformer layers by wrapping them with a small proxy that records the output and forwards all other attribute access to the original layer.

SKILL.md

3.3 KB, as published. Nobody here has run it

Capture Hidden States via Layer Hook Wrapper

When to use

You need the outputs of specific transformer layers (by index) during a forward pass.
The model does not expose output_hidden_states=True (e.g., MLX mlx_lm transformers) or you only want a few layers and cannot afford to materialize all of them.
You want to patch the model once, then read captured states off model._hidden_states after each call — no callback plumbing.

Pattern

class _LayerHook:
    """Transparent proxy: run the layer, store output, forward getattr."""
    def __init__(self, layer, idx, storage):
        self._layer, self._idx, self._storage = layer, idx, storage

    def __call__(self, *args, **kwargs):
        self._storage[self._idx] = out = self._layer(*args, **kwargs)
        return out

    def __getattr__(self, name):
        # Anything not stored on the proxy (weights, parameters()) goes to the real layer.
        return getattr(self._layer, name)


def _patch_model(model, layer_ids):
    if hasattr(model, "_hidden_states"):
        return          # idempotent
    model._hidden_states = [None] * len(layer_ids)
    layers = _get_layers(model)
    for i, lid in enumerate(layer_ids):
        layers[lid] = _LayerHook(layers[lid], i, model._hidden_states)


def _get_layers(model):
    """Duck-type well-known wrapper paths."""
    for path in (lambda m: m.model.layers,
                 lambda m: m.language_model.layers,
                 lambda m: m.layers):
        try: return path(model)
        except AttributeError: continue
    raise AttributeError(f"Cannot find layers in {type(model).__name__}")

Usage

_patch_model(target_model, draft.config.target_layer_ids)
logits = target_model(prompt[None], cache)
hidden = mx.concatenate(target_model._hidden_states, axis=-1)

Why a proxy instead of a forward hook

Works with frameworks that have no register_forward_hook (MLX).
Zero allocation cost beyond the list of outputs — the proxy is a thin wrapper.
__getattr__ delegation means layers[lid].parameters(), layers[lid].weight, etc., still work — so e.g. model.save_weights() continues to pick up the wrapped layer's parameters without knowing the wrapper exists.
Idempotency guard (hasattr(model, "_hidden_states")) makes double-patching a no-op.

Gotchas

Storage is a plain list, not thread-safe; if you call the model concurrently, use a per-call storage or lock.
The proxy sets _layer, _idx, _storage on itself; if the wrapped layer happens to have an attribute with the same name it will be shadowed. Prefix with _ to reduce collision risk.
Unpatching is not implemented; if you need to restore the layer, keep the original references you replaced and swap them back.