[bugfix] fix CI test failures for HSTUMatch

- Fix CPU CI: remove @parameterized.expand from test_hstu_match_export
  (single-case test) to fix skipIf+parameterized decorator interaction
  that prevented GPU unavailable skip
- Fix GPU CI: change candidate group from JAGGED_SEQUENCE to DEEP with
  id_feature to fix type mismatch (string vs int64) in mock data join
  during integration test. Negative sampling with standard row-append
  works correctly with DEEP candidate group.
- Update HSTUMatchItemTower to read from DEEP group key (not .sequence)
- Update _build_batch to use NEG_DATA_GROUP for candidate items

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

Author: tiankongdeguiji

tzrec/models/hstu.py

@@ -75,6 +75,7 @@ def __init__(
         contextual_group_name: str = "contextual",
     ) -> None:
         super().__init__(tower_config, output_dim, similarity, feature_group, features)
+        self._pass_grouped_features = True
         hstu_cfg = tower_config.hstu
         uih_dim = sum(feature_group_dims)
         stu_dim = hstu_cfg.stu.embedding_dim
@@ -192,6 +193,7 @@ def __init__(
         # Override _group_name: parent sets it from tower_config.input ("uih"),
         # but item tower needs to read from the candidate feature group.
         self._group_name = feature_group.group_name
+        self._pass_grouped_features = True
         cand_dim = sum(feature_group_dims)
         self._item_projection: torch.nn.Module = torch.nn.Sequential(
             torch.nn.Linear(cand_dim, output_dim),
@@ -207,7 +209,7 @@ def forward(self, grouped_features: Dict[str, torch.Tensor]) -> torch.Tensor:
         Returns:
             L2-normalized item embeddings of shape (sum_candidates, D).
         """
-        cand_emb = grouped_features[f"{self._group_name}.sequence"]
+        cand_emb = grouped_features[self._group_name]
         item_emb = self._item_projection(cand_emb)
         return F.normalize(item_emb, p=2.0, dim=-1, eps=1e-6)
 
@@ -268,7 +270,7 @@ def __init__(
         cand_features = self.get_features_in_feature_groups([cand_fg])
 
         uih_dims = self.embedding_group.group_dims(tower_cfg.input + ".sequence")
-        cand_dims = self.embedding_group.group_dims("candidate.sequence")
+        cand_dims = self.embedding_group.group_dims("candidate")
 
         # Optional contextual features
         contextual_feature_dim = 0

tzrec/models/hstu_test.py

@@ -12,7 +12,6 @@
 import unittest
 
 import torch
-from parameterized import parameterized
 from torchrec import KeyedJaggedTensor
 
 from tzrec.datasets.utils import BASE_DATA_GROUP, NEG_DATA_GROUP, Batch
@@ -43,8 +42,8 @@ def _build_model_config():
         ),
         model_pb2.FeatureGroupConfig(
             group_name="candidate",
-            feature_names=["candidate_ids"],
-            group_type=model_pb2.FeatureGroupType.JAGGED_SEQUENCE,
+            feature_names=["item_id"],
+            group_type=model_pb2.FeatureGroupType.DEEP,
         ),
     ]
     return model_pb2.ModelConfig(
@@ -94,16 +93,14 @@ def _build_features():
             )
         ),
         feature_pb2.FeatureConfig(
-            sequence_id_feature=feature_pb2.IdFeature(
-                feature_name="candidate_ids",
-                sequence_length=10,
+            id_feature=feature_pb2.IdFeature(
+                feature_name="item_id",
                 embedding_dim=48,
-                num_buckets=3953,
-                embedding_name="historical_ids",
+                num_buckets=1000,
             )
         ),
     ]
-    return create_features(feature_cfgs)
+    return create_features(feature_cfgs, neg_fields=["item_id"])
 
 
 def _build_model(device):
@@ -126,17 +123,19 @@ def _build_batch(device):
     """Build test batch with 2 users.
 
     UIH: user1 has 3 items, user2 has 4 items.
-    Candidates: 2 positive (1 per user) + 2 negative items.
+    Candidates: 2 pos (1 per user) + 2 neg items.
     """
+    # BASE: UIH sequences + positive items
     sparse_feature = KeyedJaggedTensor.from_lengths_sync(
-        keys=["historical_ids", "candidate_ids"],
-        values=torch.tensor([1, 2, 3, 4, 5, 6, 7, 10, 11, 12, 13]),
-        lengths=torch.tensor([3, 4, 2, 2]),  # uih: [3,4], candidate: [2,2]
+        keys=["historical_ids"],
+        values=torch.tensor([1, 2, 3, 4, 5, 6, 7]),
+        lengths=torch.tensor([3, 4]),
     )
+    # NEG: positive items (first batch_size) + negative items
     neg_sparse_feature = KeyedJaggedTensor.from_lengths_sync(
-        keys=["candidate_ids"],
-        values=torch.tensor([20, 21, 22, 23]),
-        lengths=torch.tensor([2, 2]),
+        keys=["item_id"],
+        values=torch.tensor([10, 11, 20, 21]),
+        lengths=torch.tensor([1, 1, 1, 1]),  # 2 pos + 2 neg, each 1 item
     )
     return Batch(
         sparse_features={
@@ -180,21 +179,20 @@ def test_hstu_match_eval(self) -> None:
         self.assertIn("recall@1", metric_result)
 
     @unittest.skipIf(*gpu_unavailable)
-    @parameterized.expand([[TestGraphType.FX_TRACE]])
-    def test_hstu_match_export(self, graph_type) -> None:
+    def test_hstu_match_export(self) -> None:
         """Test HSTUMatch export: FX trace for serving."""
         device = torch.device("cuda")
         hstu = _build_model(device)
         batch = _build_batch(device)
 
         hstu.eval()
-        hstu = create_test_model(hstu, graph_type)
+        hstu = create_test_model(hstu, TestGraphType.FX_TRACE)
         predictions = hstu(batch)
 
         self.assertIn("similarity", predictions)
         sim = predictions["similarity"]
         self.assertEqual(sim.dim(), 2)
-        self.assertEqual(sim.size(0), 2)  # batch_size
+        self.assertEqual(sim.size(0), 2)
 
     @unittest.skipIf(*gpu_unavailable)
     def test_hstu_match_predict(self) -> None:
@@ -210,7 +208,7 @@ def test_hstu_match_predict(self) -> None:
         self.assertIn("similarity", predictions)
         sim = predictions["similarity"]
         self.assertEqual(sim.dim(), 2)
-        self.assertEqual(sim.size(0), 2)  # batch_size
+        self.assertEqual(sim.size(0), 2)
         self.assertFalse(torch.isnan(sim).any())
 
 

tzrec/models/match_model.py

@@ -220,6 +220,7 @@ def __init__(
         self._similarity = similarity
         self._feature_group = feature_group
         self._features = features
+        self._pass_grouped_features = False
 
 
 class MatchModel(BaseModel):
@@ -492,8 +493,9 @@ def predict(self, batch: Batch) -> Dict[str, torch.Tensor]:
             embedding (dict): tower output embedding.
         """
         grouped_features = self.embedding_group(batch)
-        return {
-            f"{self._tower_name}_emb": getattr(self, self._tower_name)(
-                grouped_features[self._group_name]
-            )
-        }
+        tower = getattr(self, self._tower_name)
+        if tower._pass_grouped_features:
+            tower_input = grouped_features
+        else:
+            tower_input = grouped_features[self._group_name]
+        return {f"{self._tower_name}_emb": tower(tower_input)}

tzrec/tests/configs/hstu_fg_mock.config

@@ -59,11 +59,9 @@ feature_configs {
     }
 }
 feature_configs {
-    sequence_id_feature {
+    id_feature {
         feature_name: "item_id"
         expression: "item:item_id"
-        sequence_length: 10
-        sequence_delim: ";"
         num_buckets: 1000
         embedding_dim: 48
     }
@@ -78,7 +76,7 @@ model_config {
     feature_groups {
         group_name: "candidate"
         feature_names: "item_id"
-        group_type: JAGGED_SEQUENCE
+        group_type: DEEP
     }
     hstu_match {
         hstu_tower {