[Refactor] msgspec standardisation for data entry key names and improved type checks (#3149)

Signed-off-by: Divyansh Singhvi <divyanshsinghvi@gmail.com>
Signed-off-by: dsinghvi <divyanshsinghvi@gmail.com>

tests/model_executor/models/cosyvoice3/test_cosyvoice3_model_helpers.py

@@ -164,11 +164,12 @@ def test_forward_prefers_token_offset_when_present():
 
     runtime_info = [
         {
-            "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
-            "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
-            "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
-            "token_offset": 2,
-            "left_context_size": 1,
+            "embed": {
+                "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+                "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+                "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            },
+            "meta": {"left_context_size": 2},
         }
     ]
 
@@ -193,10 +194,12 @@ def test_forward_falls_back_to_left_context_size_for_backward_compat():
 
     runtime_info = [
         {
-            "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
-            "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
-            "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
-            "left_context_size": 2,
+            "embed": {
+                "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+                "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+                "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            },
+            "meta": {"left_context_size": 2},
         }
     ]
 
@@ -214,10 +217,12 @@ def test_forward_ignores_single_request_padded_tail_tokens():
     model = _make_code2wav_model(with_stride_cfg=True)
     runtime_info = [
         {
-            "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
-            "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
-            "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
-            "token_offset": 0,
+            "embed": {
+                "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+                "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+                "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            },
+            "meta": {"left_context_size": 0},
         }
     ]
 
@@ -238,10 +243,12 @@ def test_forward_uses_non_stream_decode_without_chunk_metadata():
 
     runtime_info = [
         {
-            "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
-            "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
-            "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
-            "prefix_ids": [101, 102],
+            "embed": {
+                "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+                "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+                "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            },
+            "ids": {"prompt": [101, 102]},
             "generated_len": 3,
         }
     ]
@@ -275,12 +282,16 @@ def test_forward_reuses_streaming_cache_state_between_chunks():
     )
     runtime_info = [
         {
-            "req_id": ["rid-stream"],
-            "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
-            "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
-            "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
-            "token_offset": 0,
-            "stream_finished": torch.tensor(False),
+            "embed": {
+                "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+                "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+                "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            },
+            "meta": {
+                "req_id": ["rid-stream"],
+                "stream_finished": torch.tensor(False),
+                "left_context_size": 0,
+            },
         }
     ]
 
@@ -321,12 +332,16 @@ def test_forward_clears_streaming_cache_on_terminal_chunk():
     )
     runtime_info = [
         {
-            "req_id": ["rid-stream"],
-            "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
-            "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
-            "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
-            "token_offset": 0,
-            "stream_finished": torch.tensor(False),
+            "embed": {
+                "speech_token": torch.tensor([[1, 2, 3]], dtype=torch.long),
+                "speech_feat": torch.tensor([[[0.1, 0.2], [0.3, 0.4]]], dtype=torch.float32),
+                "embedding": torch.tensor([[0.5, 0.6]], dtype=torch.float32),
+            },
+            "meta": {
+                "req_id": ["rid-stream"],
+                "stream_finished": torch.tensor(False),
+                "left_context_size": 0,
+            },
         }
     ]
 
@@ -338,7 +353,7 @@ def test_forward_clears_streaming_cache_on_terminal_chunk():
     )
     assert "rid-stream" in model._stream_vocoder_cache_by_req
 
-    runtime_info[0]["stream_finished"] = torch.tensor(True)
+    runtime_info[0]["meta"]["stream_finished"] = torch.tensor(True)
     out = model.forward(
         input_ids=torch.tensor([0, 1, 2], dtype=torch.long),
         positions=torch.tensor([0, 1, 2], dtype=torch.long),

tests/model_executor/stage_input_processors/test_cosyvoice3_stage_input_processors.py

@@ -64,9 +64,11 @@ def test_talker2code2wav_async_chunk_final_payload_uses_absolute_token_offset():
         external_req_id="rid-0",
         output_token_ids=[1, 2, 6562, 3],
         additional_information={
-            "speech_token": [torch.tensor([[11, 12, 13]])],
-            "speech_feat": [torch.tensor([[[0.1, 0.2], [0.3, 0.4]]])],
-            "embedding": [torch.tensor([[0.5, 0.6]])],
+            "embed": {
+                "speech_token": [torch.tensor([[11, 12, 13]])],
+                "speech_feat": [torch.tensor([[[0.1, 0.2], [0.3, 0.4]]])],
+                "embedding": [torch.tensor([[0.5, 0.6]])],
+            },
         },
         is_finished=lambda: True,
     )
@@ -79,15 +81,14 @@ def test_talker2code2wav_async_chunk_final_payload_uses_absolute_token_offset():
     )
 
     assert payload is not None
-    assert payload["meta"]["finished"].item() is True
-    assert payload["codes"]["audio"] == [1, 2, 3]
-    assert payload["token_offset"] == 0
-    assert payload["left_context_size"] == 0
-    assert payload["req_id"] == ["rid-0"]
-    assert payload["stream_finished"].item() is True
-    assert "speech_token" in payload
-    assert "speech_feat" in payload
-    assert "embedding" in payload
+    assert payload.meta.finished.item() is True
+    assert payload.codes.audio.tolist() == [1, 2, 3]
+    assert payload.meta.left_context_size == 0
+    assert payload.meta.req_id == ["rid-0"]
+    assert payload.meta.stream_finished.item() is True
+    assert payload.embed.speech_token is not None
+    assert payload.embed.speech_feat is not None
+    assert payload.embed.embedding is not None
 
 
 def test_talker2code2wav_async_chunk_emits_eof_when_finished_without_valid_codes():
@@ -107,8 +108,8 @@ def test_talker2code2wav_async_chunk_emits_eof_when_finished_without_valid_codes
     )
 
     assert payload is not None
-    assert payload["codes"]["audio"] == []
-    assert payload["meta"]["finished"].item() is True
+    assert payload.codes.audio.tolist() == []
+    assert payload.meta.finished.item() is True
 
 
 def test_talker2code2wav_async_chunk_does_not_reemit_without_new_tokens():
@@ -134,8 +135,8 @@ def test_talker2code2wav_async_chunk_does_not_reemit_without_new_tokens():
     )
 
     assert payload1 is not None
-    assert payload1["codes"]["audio"] == [1, 2]
-    assert payload1["token_offset"] == 0
+    assert payload1.codes.audio.tolist() == [1, 2]
+    assert payload1.meta.left_context_size == 0
     assert payload2 is None
 
 
@@ -164,9 +165,9 @@ def test_talker2code2wav_async_chunk_waits_for_prelookahead_and_emits_cumulative
 
     assert payload_pending is None
     assert payload_ready is not None
-    assert payload_ready["codes"]["audio"] == [1, 2, 3]
-    assert payload_ready["token_offset"] == 0
-    assert payload_ready["meta"]["finished"].item() is False
+    assert payload_ready.codes.audio.tolist() == [1, 2, 3]
+    assert payload_ready.meta.left_context_size == 0
+    assert payload_ready.meta.finished.item() is False
 
 
 def test_talker2code2wav_async_chunk_final_flush_uses_previous_token_offset():
@@ -193,13 +194,13 @@ def test_talker2code2wav_async_chunk_final_flush_uses_previous_token_offset():
     )
 
     assert payload_stream is not None
-    assert payload_stream["meta"]["finished"].item() is False
-    assert payload_stream["codes"]["audio"] == [3, 4, 5]
-    assert payload_stream["token_offset"] == 0
+    assert payload_stream.meta.finished.item() is False
+    assert payload_stream.codes.audio.tolist() == [3, 4, 5]
+    assert payload_stream.meta.left_context_size == 0
     assert payload_final is not None
-    assert payload_final["meta"]["finished"].item() is True
-    assert payload_final["codes"]["audio"] == [3, 4, 5, 6]
-    assert payload_final["token_offset"] == 2
+    assert payload_final.meta.finished.item() is True
+    assert payload_final.codes.audio.tolist() == [3, 4, 5, 6]
+    assert payload_final.meta.left_context_size == 2
 
 
 def test_talker2code2wav_async_chunk_respects_prompt_token_pad_on_first_chunk():
@@ -208,7 +209,7 @@ def test_talker2code2wav_async_chunk_respects_prompt_token_pad_on_first_chunk():
         external_req_id="rid-pad",
         output_token_ids=[8, 9, 10],
         additional_information={
-            "speech_token": [torch.tensor([[1, 2, 3]])],
+            "embed": {"speech_token": [torch.tensor([[1, 2, 3]])]},
         },
         is_finished=lambda: False,
     )
@@ -229,8 +230,8 @@ def test_talker2code2wav_async_chunk_respects_prompt_token_pad_on_first_chunk():
 
     assert payload_pending is None
     assert payload_ready is not None
-    assert payload_ready["codes"]["audio"] == [8, 9, 10, 11]
-    assert payload_ready["token_offset"] == 0
+    assert payload_ready.codes.audio.tolist() == [8, 9, 10, 11]
+    assert payload_ready.meta.left_context_size == 0
 
 
 def test_talker2code2wav_async_chunk_emits_terminal_eof_without_duplicate_audio():
@@ -256,8 +257,8 @@ def test_talker2code2wav_async_chunk_emits_terminal_eof_without_duplicate_audio(
     )
 
     assert payload_stream is not None
-    assert payload_stream["meta"]["finished"].item() is False
-    assert payload_stream["codes"]["audio"] == [3, 4]
+    assert payload_stream.meta.finished.item() is False
+    assert payload_stream.codes.audio.tolist() == [3, 4]
     assert payload_final is not None
-    assert payload_final["meta"]["finished"].item() is True
-    assert payload_final["codes"]["audio"] == []
+    assert payload_final.meta.finished.item() is True
+    assert payload_final.codes.audio.tolist() == []

tests/model_executor/stage_input_processors/test_mimo_audio_flush_remaining_codes.py

@@ -20,16 +20,16 @@ def test_flush_remaining_codes_when_no_codes_accumulated_missing_request_id():
     """No entry for request_id: treat as empty, return finished sentinel with empty codes."""
     tm = SimpleNamespace(code_prompt_token_ids={})
     out = _flush_remaining_codes(tm, "missing", chunk_size=3, left_context_size=3)
-    assert out["codes"]["audio"] == _sentinel()["codes"]["audio"]
-    assert out["meta"]["finished"].item() is True
+    assert out.codes.audio.tolist() == _sentinel()["codes"]["audio"]
+    assert out.meta.finished.item() is True
 
 
 def test_flush_remaining_codes_when_no_codes_accumulated_empty_list():
     """Explicit empty accumulation list returns the same sentinel."""
     tm = SimpleNamespace(code_prompt_token_ids={"r": []})
     out = _flush_remaining_codes(tm, "r", chunk_size=3, left_context_size=3)
-    assert out["codes"]["audio"] == []
-    assert out["meta"]["finished"].item() is True
+    assert out.codes.audio.tolist() == []
+    assert out.meta.finished.item() is True
 
 
 def test_flush_remaining_codes_partial_chunk_remaining():
@@ -41,8 +41,8 @@ def test_flush_remaining_codes_partial_chunk_remaining():
         code_prompt_token_ids={"r": [[1], [2], [3], [4], [5], [6], [7]]},
     )
     out = _flush_remaining_codes(tm, "r", chunk_size=3, left_context_size=3)
-    assert out["meta"]["finished"].item() is True
-    assert out["codes"]["audio"] == [4, 5, 6, 7]
+    assert out.meta.finished.item() is True
+    assert out.codes.audio.tolist() == [4, 5, 6, 7]
 
 
 def test_flush_remaining_codes_when_length_is_exact_multiple_of_chunk_size():
@@ -52,7 +52,7 @@ def test_flush_remaining_codes_when_length_is_exact_multiple_of_chunk_size():
     )
     out = _flush_remaining_codes(tm, "r", chunk_size=3, left_context_size=3)
     # context_length = chunk_size = 3, end_index = min(6, 6) -> all 6
-    assert out["codes"]["audio"] == [1, 2, 3, 4, 5, 6]
+    assert out.codes.audio.tolist() == [1, 2, 3, 4, 5, 6]
 
 
 @pytest.mark.parametrize(
@@ -74,5 +74,5 @@ def test_flush_remaining_codes_context_window_end_index(
     tm = SimpleNamespace(code_prompt_token_ids={"r": accumulated})
     out = _flush_remaining_codes(tm, "r", chunk_size=chunk_size, left_context_size=left_context)
     expected_flat = list(range(length - expected_end_index, length))
-    assert out["codes"]["audio"] == expected_flat
-    assert out["meta"]["finished"].item() is True
+    assert out.codes.audio.tolist() == expected_flat
+    assert out.meta.finished.item() is True

tests/model_executor/stage_input_processors/test_qwen3_tts_async_chunk.py

@@ -80,8 +80,8 @@ def test_eof_marker_when_finished_empty():
         request=_req("r", finished=True),
         is_finished=True,
     )
-    assert p["codes"] == {"audio": []}
-    assert p["meta"]["finished"].item() is True
+    assert p.codes.audio.tolist() == []
+    assert p.meta.finished.item() is True
 
 
 def test_flush_on_finish():
@@ -94,8 +94,8 @@ def test_flush_on_finish():
         is_finished=True,
     )
     assert p is not None
-    assert p["meta"]["finished"].item() is True
-    assert len(p["codes"]["audio"]) == _Q * 24
+    assert p.meta.finished.item() is True
+    assert len(p.codes.audio) == _Q * 24
 
 
 _CASES = [
@@ -149,8 +149,8 @@ def test_streaming_phases(config, n_frames, finished, expected):
     else:
         exp_ctx, exp_window = expected
         assert payload is not None
-        assert payload["meta"]["left_context_size"] == exp_ctx
-        assert len(payload["codes"]["audio"]) == _Q * exp_window
+        assert payload.meta.left_context_size == exp_ctx
+        assert len(payload.codes.audio) == _Q * exp_window
 
 
 def test_dynamic_ic_adapts_to_load():
@@ -160,14 +160,14 @@ def test_dynamic_ic_adapts_to_load():
     # Low load (1/8) -> IC=2 -> emit at 2
     p1 = _call(tm, "r", n_frames=2)
     assert p1 is not None
-    assert len(p1["codes"]["audio"]) == _Q * 2
+    assert len(p1.codes.audio) == _Q * 2
 
     # High load on a new request: active=6/8 -> IC=8 -> emit at 8
     for i in range(4):
         tm.code_prompt_token_ids[f"other-{i}"] = [[0]]
     p2 = _call(tm, "new-high-load", n_frames=8)
     assert p2 is not None
-    assert len(p2["codes"]["audio"]) == _Q * 8
+    assert len(p2.codes.audio) == _Q * 8
 
     # Requests past initial phase still count in load factor
     tm2 = _tm(max_num_seqs=4)
@@ -176,7 +176,7 @@ def test_dynamic_ic_adapts_to_load():
     # active=4/4=1.0 -> IC=16
     p3 = _call(tm2, "new", n_frames=16)
     assert p3 is not None
-    assert len(p3["codes"]["audio"]) == _Q * 16
+    assert len(p3.codes.audio) == _Q * 16
 
 
 def test_ic_load_change_mid_request():
@@ -195,7 +195,7 @@ def test_ic_load_change_mid_request():
     assert _call(tm, "r", n_frames=25) is None
     p3 = _call(tm, "r", n_frames=27)
     assert p3 is not None
-    assert p3["meta"]["left_context_size"] == 2
+    assert p3.meta.left_context_size == 2
 
     # A *new* request under high load gets IC=16 (not IC=2).
     # Frame 2 would emit under IC=2 but must hold under IC=16.
@@ -214,13 +214,13 @@ def test_connector_initial_chunk_config_overrides_dynamic_ic():
 
     p1 = _call(tm, "r", n_frames=4)
     assert p1 is not None
-    assert len(p1["codes"]["audio"]) == _Q * 4
+    assert len(p1.codes.audio) == _Q * 4
 
     # Only the first chunk uses the small size; the next emit is 4+25.
     assert _call(tm, "r", n_frames=25) is None
     p2 = _call(tm, "r", n_frames=29)
     assert p2 is not None
-    assert p2["meta"]["left_context_size"] == 4
+    assert p2.meta.left_context_size == 4
 
 
 @pytest.mark.parametrize(
@@ -268,8 +268,8 @@ def test_first_streaming_chunk_prepends_ref_code_context():
     )
 
     assert payload is not None
-    assert payload["meta"]["left_context_size"] == 2
-    assert len(payload["codes"]["audio"]) == _Q * 12
+    assert payload.meta.left_context_size == 2
+    assert len(payload.codes.audio) == _Q * 12
 
 
 def test_ref_code_context_applies_to_all_streaming_chunks():
@@ -290,8 +290,8 @@ def test_ref_code_context_applies_to_all_streaming_chunks():
 
     assert payload is not None
     # ref_code (2 frames) prepended as left context on second chunk too
-    assert payload["meta"]["left_context_size"] == 10 + 2
-    assert len(payload["codes"]["audio"]) == _Q * (35 + 2)
+    assert payload.meta.left_context_size == 10 + 2
+    assert len(payload.codes.audio) == _Q * (35 + 2)
 
 
 def test_ref_code_context_can_be_buffered_before_first_emit():
@@ -325,8 +325,8 @@ def test_ref_code_context_can_be_buffered_before_first_emit():
 
     assert payload is not None
     # ref_code (2 frames) is kept (not popped) for subsequent chunks
-    assert payload["meta"]["left_context_size"] == 2
-    assert len(payload["codes"]["audio"]) == _Q * 12
+    assert payload.meta.left_context_size == 2
+    assert len(payload.codes.audio) == _Q * 12
     assert rid in tm.request_payload
 
 
@@ -353,7 +353,7 @@ def test_non_async_processor_prepends_ref_code_and_sets_trim_context():
 
     assert len(prompts) == 1
     prompt = prompts[0]
-    assert prompt["additional_information"] == {"meta": {"left_context_size": [2]}}
+    assert prompt["additional_information"] == {"meta": {"left_context_size": 2}}
     assert prompt["prompt_token_ids"] == [
         9,
         8,
@@ -401,4 +401,4 @@ def test_non_async_processor_filters_out_of_range_codec_values():
     prompt = prompts[0]
     # Only ref_code (1 frame) + 2 valid frames = 3 frames * 4 quantizers = 12 codes
     assert len(prompt["prompt_token_ids"]) == 4 * 3
-    assert prompt["additional_information"] == {"meta": {"left_context_size": [1]}}
+    assert prompt["additional_information"] == {"meta": {"left_context_size": 1}}

tests/model_executor/stage_input_processors/test_voxcpm_async_chunk.py

@@ -57,8 +57,8 @@ def test_latent2vae_async_chunk_serializes_latent_payload():
     )
 
     assert payload is not None
-    assert torch.equal(payload["finished"], torch.tensor(False, dtype=torch.bool))
-    recovered = _decode_serialized_latent(payload["code_predictor_codes"])
+    assert torch.equal(payload.meta.finished, torch.tensor(False, dtype=torch.bool))
+    recovered = _decode_serialized_latent(payload.codes.audio.tolist())
     torch.testing.assert_close(recovered, latent.to(torch.bfloat16).to(torch.float32).unsqueeze(0))
 
 
@@ -70,10 +70,8 @@ def test_latent2vae_async_chunk_returns_terminal_marker_without_latent():
         is_finished=False,
     )
 
-    assert payload == {
-        "code_predictor_codes": [],
-        "finished": torch.tensor(True, dtype=torch.bool),
-    }
+    assert payload.codes.audio.tolist() == []
+    assert torch.equal(payload.meta.finished, torch.tensor(True, dtype=torch.bool))
 
 
 def test_latent2vae_async_chunk_returns_none_for_nonterminal_empty_chunk():

tests/model_executor/stage_input_processors/test_voxtral_tts_async_chunk.py

@@ -114,8 +114,8 @@ def test_flush_tail_when_finished():
     )
 
     assert payload is not None
-    assert payload["meta"]["finished"].item() is True
-    codes = payload["codes"]["audio"]
+    assert payload.meta.finished.item() is True
+    codes = payload.codes.audio
     # Format: [ctx_frames, context_length, ...flat_codes]
     assert len(codes) >= 2  # At least ctx_frames + context_length header
     ctx_frames = codes[0]
@@ -137,8 +137,8 @@ def test_eof_marker_when_finished_with_no_frames():
         request=request,
     )
 
-    assert payload["codes"] == {"audio": []}
-    assert payload["meta"]["finished"].item() is True
+    assert payload.codes.audio.tolist() == []
+    assert payload.meta.finished.item() is True
 
 
 def test_normal_chunk_emission():
@@ -162,7 +162,7 @@ def test_normal_chunk_emission():
 
     # A chunk should be emitted
     assert payload is not None
-    codes = payload["codes"]["audio"]
+    codes = payload.codes.audio
     ctx_frames = codes[0]
     context_length = codes[1]
     assert ctx_frames == 20  # 25 - 5(chunk_size_at_begin)
@@ -189,7 +189,7 @@ def test_small_initial_chunks():
         )
 
     assert payload is not None
-    codes = payload["codes"]["audio"]
+    codes = payload.codes.audio
     ctx_frames = codes[0]
     context_length = codes[1]
     assert ctx_frames == 0
@@ -238,19 +238,16 @@ def test_context_handling_format():
         )
 
     assert payload is not None
-    codes = payload["codes"]["audio"]
-    # First two elements are ctx_frames and context_length
-    ctx_frames = codes[0]
-    context_length = codes[1]
+    codes = payload.codes.audio
+    # codes is a 1-D long tensor: [ctx_frames, context_length, ...flat_codes]
+    ctx_frames = int(codes[0].item())
+    context_length = int(codes[1].item())
     flat_codes = codes[2:]
-    # The window has ctx_frames + context_length frames
-    assert isinstance(ctx_frames, int)
-    assert isinstance(context_length, int)
     assert ctx_frames >= 0
     assert context_length > 0
     # flat_codes = total_window_frames * codebook_dim
     total_window_frames = ctx_frames + context_length
-    assert len(flat_codes) == total_window_frames * 2  # codebook_dim=2
+    assert flat_codes.numel() == total_window_frames * 2  # codebook_dim=2
 
 
 def test_none_pooling_output_not_finished_returns_none():