[bugfix] fix contextual_seq_len not passed from preprocessor to STULayer

[tiankongdeguiji]

Summary

• contextual_seq_len in STULayer was configured from the STU proto field (defaulting to 0), instead of being derived from InputPreprocessor.contextual_seq_len(). This caused incorrect attention masks in both Triton and PyTorch kernels when contextual features were used but contextual_seq_len was not explicitly set in the proto config.
• Fixed by reordering HSTUTransducer.__init__ to create InputPreprocessor first, then override stu["contextual_seq_len"] with the computed value before constructing STULayer instances.
• Marked the contextual_seq_len field in STU proto as deprecated (kept for backward compatibility).
• Bumped version to 1.1.3.

Test plan

• python -m tzrec.modules.gr.stu_test passes
• pre-commit run -a passes
• Verify attention mask correctness with contextual features enabled

🤖 Generated with Claude Code

[github-actions]

Code Review Summary

Clean, well-targeted bugfix. The root cause is clear and the fix is minimal — reordering __init__ so InputPreprocessor is created before STULayer instances, allowing contextual_seq_len to be derived correctly. The proto deprecation comment is accurate and the version bump is appropriate.

Inline Comments

1. Defensive dict copy (hstu_transducer.py:82) — The stu dict is mutated in-place, which works today but is fragile. Suggest stu = dict(stu) before the assignment.
2. Runtime deprecation warning (module.proto:258) — Consider logging a warning when a user has explicitly set the now-ignored contextual_seq_len proto field.

Test Coverage Gap

There are no tests for HSTUTransducer anywhere in the test suite. The existing stu_test.py tests STULayer directly by passing contextual_seq_len as a constructor argument — it never exercises the config-to-module wiring where this bug lived. A minimal integration test that constructs HSTUTransducer with contextual features and asserts _stu_module layers have the correct _contextual_seq_len would guard against regression.

Pre-existing Doc Issue (not from this PR)

In tzrec/modules/gr/stu.py:190, the contextual_seq_len docstring says (bool) but the actual type is int. Worth fixing while in the neighborhood.

---

Overall: Approve with minor suggestions. The fix is correct and important — the Triton kernel was being compiled with HAS_CONTEXTUAL_SEQ_LEN=False when contextual features were present, producing wrong attention masks.

🤖 Generated with Claude Code

[github-actions]

Code Review Summary

Clean, well-scoped bugfix. The root cause (proto default silently overriding the preprocessor-derived value) is correctly addressed, and the defensive shallow copy + deprecation warning are good additions. The docstring type fix (bool → int) and proto comment update are both accurate.

Noteworthy Feedback

1. Consider a regression test for contextual_seq_len propagation (test coverage)

No existing test verifies that HSTUTransducer correctly passes contextual_seq_len from the preprocessor to the STULayer instances. The existing stu_test.py constructs STULayer directly (bypassing the transducer), and dlrm_hstu_test.py only checks output shapes. A simple CPU-only unit test would guard this fix:

# Construct HSTUTransducer with contextual features, then assert:
for layer in transducer._stu_module._stu_layers:
    assert layer._contextual_seq_len == transducer._input_preprocessor.contextual_seq_len()

2. Use proto [deprecated = true] annotation (see inline comment on module.proto)

3. Minor: warnings.warn vs logger.warning (see inline comment on hstu_transducer.py)

No Issues Found

• Performance: no concerns — changes are init-only, no hot-path impact.
• Security: no concerns — internal config plumbing, no external input.

LGTM with the test suggestion above as a nice-to-have.

🤖 Generated with Claude Code