[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

@@ -207,7 +207,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 +268,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/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 {