[flyingdoubleG]

Summary

CrossLayerTranscoder.encode_single_layer is used by FeatureAnalyzer for CLT feature analysis, after keep_only_decoders_for_layer_from. With sparsity_include_decoder_norm=True (the default, also used by examples/train_pythia_clt_topk.py), it disagrees with encode in two ways.

1. Wrong norm tensor. It multiplies the (..., d_sae) pre-activations by decoder_norm_per_feature(layer=layer). That is an (n_layers, d_sae) table of per-target norms, with zero rows before layer.
   • With a (batch, d_model) input this raises RuntimeError: The size of tensor a (5) must match the size of tensor b (3).
   • When a leading dimension happens to equal n_layers (e.g. seq_len == n_layers), it broadcasts silently and gives wrong activations.
2. No rescale after the activation. It never divides the activations back by the norm, which encode does.

Reproduction

encode_single_layer(x[:, l], l) compared with encode(x)[:, l] on a 3-layer CLT:

before: layer 0..2: RuntimeError: The size of tensor a (5) must match the size of tensor b (3) at non-singleton dimension 0
        seq_len == n_layers: max |diff| = 2.418
after:  layer 0..2: max |diff| = 0.0000;  seq_len == n_layers: max |diff| = 2.4e-07

Change

• Use the layer's decoder norm concatenated over its target layers. This is row layer of decoder_norm_per_feature(), which encode uses, computed as the root sum of squares of the per-target norms, so it also works with the decoders kept by keep_only_decoders_for_layer_from.
• Divide the activations by that norm after the activation function.

Tests

• New tests/unit/test_clt.py checks that encode_single_layer equals the corresponding slice of encode, with and without sparsity_include_decoder_norm and with and without keep_only_decoders_for_layer_from.
• The sparsity_include_decoder_norm=True cases fail before this change and pass after.
• ruff check and ruff format --check are clean.