Skip to content

docs/tests: feature_intervention_generate logits changed from 3D to 2D - document / test this#90

Open
hijohnnylin wants to merge 1 commit intodecoderesearch:mainfrom
hijohnnylin:logits_dims
Open

docs/tests: feature_intervention_generate logits changed from 3D to 2D - document / test this#90
hijohnnylin wants to merge 1 commit intodecoderesearch:mainfrom
hijohnnylin:logits_dims

Conversation

@hijohnnylin
Copy link
Copy Markdown
Collaborator

@hijohnnylin hijohnnylin commented Mar 24, 2026

When adding support for nnsight engine, we slightly changed the return shape of feature_intervention_generate's logits tensor. It changed from batch, seq_len, vocab_size to just seq_len, vocab_size. This broke downstream code/users that assumed the previous shape.

Previous: No squeeze(0), returned logits is 2D

return generation, logits, activations

New: squeeze(0). returned logits is 3D

return generation, logits.squeeze(0), activations

The new shape is fine since the batch size is always 1 anyway, so this commit is a Claude-assisted update that doesn't update the behavior - it only adds documentation about this as well as a few tests and new asserts.
Possibly it's too verbose a PR for the relatively small change, so feel free to edit/change/minimize it.

(Downstream we added a check to ensure it's dimension 2 which is our workaround for now.)

@hijohnnylin hijohnnylin requested a review from hannamw March 24, 2026 00:04
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant