Fix beam search test helper indexing

[fallintoplace]

What changed

This fixes _check_content_equal_per_beam() in tests/generate/beam_search_test.py so it iterates by batch_idx and beam_idx when comparing replicated beam content back to the original batch entries.

It also adds a focused regression test that exercises both 2D beam-expanded content and 1D beam-expanded flags.

Why

The helper previously stepped i by beam_size and then indexed beam_content[i * beam_size] and original_content[i].

With batch_size=2 and beam_size=2, the second loop iteration reached beam_content[4] and original_content[2], which are both out of bounds.

Impact

The beam search tests now validate per-batch/per-beam replication correctly, and the regression test covers the path that used to go out of bounds.

Validation

• uv run --extra test python tests/generate/beam_search_test.py