[refactor] Refactor HSTUMatch to use STU module with jagged sequences

[tiankongdeguiji]

Summary

• Replace old HSTUEncoder/SequentialTransductionUnitJagged with modern STUStack/STULayer from gr/stu.py in HSTUMatch model
• Add new SequencePreprocessor for match models — handles UIH-only sequences with optional contextual features, timestamps, and actions (no candidate concatenation)
• Switch to JAGGED_SEQUENCE feature groups for native jagged tensor support with separate UIH and candidate groups
• Update proto definitions: HSTUMatchTower uses HSTU config, new GRSequencePreprocessor in GRInputPreprocessor oneof
• Fix WhichOneof field name bug in create_output_postprocessor
• Remove unused HSTUEncoder class and tests

Changes

• tzrec/protos/module.proto — Add GRSequencePreprocessor, add to GRInputPreprocessor oneof
• tzrec/protos/tower.proto — HSTUMatchTower uses HSTU + max_seq_len instead of HSTUEncoder
• tzrec/protos/models/match_model.proto — Remove output_dim from HSTUMatch
• tzrec/modules/gr/preprocessors.py — Add SequencePreprocessor class + register in factory
• tzrec/modules/gr/postprocessors.py — Fix WhichOneof("input_preprocessor") → WhichOneof("output_postprocessor")
• tzrec/models/hstu.py — Rewrite with STUStack, SequencePreprocessor, jagged two-tower architecture
• tzrec/modules/sequence.py — Remove HSTUEncoder (replaced by SequencePreprocessor + STUStack)
• tzrec/modules/sequence_test.py — Remove HSTUEncoderTest
• tzrec/models/hstu_test.py — Updated test with JAGGED_SEQUENCE groups and new proto config

Test plan

• python -m tzrec.models.hstu_test — NORMAL and FX_TRACE tests pass
• python -m tzrec.models.dlrm_hstu_test — DlrmHSTU regression test passes
• python -m tzrec.modules.sequence_test — Sequence module tests pass
• pre-commit run -a — All linting/formatting checks pass

🤖 Generated with Claude Code

[github-actions]

Bug risk: missing input validation. If len(neg_data) != len(pos_data) * neg_sample_num (e.g., sampler returns unexpected count), the offset array silently produces wrong results — corrupt training data with no error signal.

Suggested change

	def combine_neg_as_candidate_sequence(
	def combine_neg_as_candidate_sequence(
	pos_data: pa.Array,
	neg_data: pa.Array,
	neg_sample_num: int,
	seq_delim: str,
	) -> pa.Array:
	"""Combine positive and negative items into candidate sequences.
	For each sample, joins the positive item with its negative items using the
	sequence delimiter. Used when candidate features are sequence_id_features
	in a JAGGED_SEQUENCE group.
	Args:
	pos_data: positive item IDs, one per sample. Shape: (B,).
	neg_data: negative item IDs. Shape: (B * neg_sample_num,).
	neg_sample_num: number of negative samples per positive.
	seq_delim: delimiter for joining items into a sequence string.
	Returns:
	pa.Array of strings, each containing "pos;neg1;neg2;..." per sample.
	Example:
	pos_data = ["1", "2"]
	neg_data = ["3", "4", "5", "6"]
	neg_sample_num = 2
	seq_delim = ";"
	result = ["1;3;4", "2;5;6"]
	"""
	assert len(neg_data) == len(pos_data) * neg_sample_num, (
	f"neg_data length ({len(neg_data)}) must equal "
	f"pos_data length ({len(pos_data)}) * neg_sample_num ({neg_sample_num})"
	)

[github-actions]

Inaccurate docstring. The class docstring says "During training, returns one embedding per UIH position (autoregressive)", but the forward method at line 164 always extracts the last-position embedding via user_emb[seq_offsets[1:] - 1] regardless of training mode. The return shape is always (B, D).

Consider simplifying to: "Produces user embeddings by extracting the last-position embedding per user."

[github-actions]

Nit: this says "and L2 normalizes" unconditionally, but normalization only happens for COSINE similarity (line 215). Consider: "Projects candidate embeddings to STU embedding dimension. L2 normalizes when using COSINE similarity."

[github-actions]

Design note: This boolean flag toggles between two fundamentally different calling conventions (Dict[str, Tensor] vs single Tensor). It works but is type-unsafe — tower_input is implicitly Union[Dict, Tensor]. Consider whether overriding predict in subclasses or accepting group_name: Optional[str] would be cleaner long-term. Not blocking, but worth noting for maintainability.

[github-actions]

Performance: This allocates a zeros tensor on GPU every forward call, and when max_contextual_seq_len > 0, a second zeros tensor is also allocated via _fx_timestamp_contextual_zeros. Since timestamps are always zero for this preprocessor, consider pre-allocating a buffer in __init__ (via register_buffer) and slicing it, or checking whether downstream consumers (e.g., L2NormPostprocessor) can skip timestamp processing entirely when unused.

[github-actions]

Code Review Summary

Well-structured refactor that cleanly replaces the old HSTUEncoder with modern STUStack/STULayer, adds UIHPreprocessor for match models, and simplifies negative sampling. The two-tower separation is clear, proto definitions are well-organized, and the WhichOneof bugfix is a good catch. Overall a solid improvement.

Issues to Address

Correctness (2 items — see inline comments):

• combine_neg_as_candidate_sequence has no validation that len(neg_data) == len(pos_data) * neg_sample_num. Mismatches silently corrupt training data.
• user_emb[seq_offsets[1:] - 1] produces index -1 (wrapping to last element) if any user has zero-length UIH sequence — silently returns wrong embedding.

Test Coverage (2 items):

• UIHPreprocessor is ~220 lines of new code with conditional branches (contextual features, action encoder) but has no unit test. It's only exercised indirectly through integration tests. A dedicated test in preprocessors_test.py (following the existing hypothesis-based pattern) would be high-value.
• All hstu_test.py tests are GPU-only (@unittest.skipIf(*gpu_unavailable)). CPU-only CI gets zero coverage for the entire HSTUMatch model. At minimum, a test verifying model instantiation and config validation without a forward pass could run on CPU.

Documentation (2 items — see inline comments):

• HSTUMatchUserTower docstring claims different train/inference behavior, but forward always extracts last-position embedding.
• HSTUMatchItemTower docstring says "L2 normalizes" unconditionally, but it's conditional on COSINE similarity.

Design Notes (non-blocking)

• The _pass_grouped_features boolean dispatch in match_model.py toggles between Dict and Tensor calling conventions — pragmatic but type-unsafe. See inline comment.
• HSTUMatchItemTower overrides _group_name immediately after the parent sets it — a "set then override" pattern that could be cleaner with an explicit parameter.
• UIHPreprocessor and ContextualInterleavePreprocessor share substantial duplicated logic (contextual projection, concat_2D_jagged, action encoder setup). Extracting shared methods would reduce maintenance burden.
• Hardcoded zero timestamps allocated every forward call in UIHPreprocessor — could be pre-allocated or skipped when downstream postprocessors don't use them.