Add validation logic for keys & weights of KJT

Summary:
Introduces validations for the correctness of `keys` and `weights` within a `KeyedJaggedTensor` (KJT), ensuring that KJTs are properly formed before they are used in downstream computations.

Added validation rules:

 - `keys` should NOT contain duplications
 - For non-VBE scenario, `lengths` size should be divisible by the number of keys.
 - If `keys` is empty, the KJT should be an empty one.
 - `weights` must match the size of `values`.

VBE scenario will be addressed in follow-up changes.

Reviewed By: TroyGarden

Differential Revision: D77247081

fbshipit-source-id: bf0a9915a5a499e5504c912c5d720d8e80928026

Author: Shuangping Liu

torchrec/sparse/jagged_tensor_validator.py

@@ -17,18 +17,21 @@ def validate_keyed_jagged_tensor(
     """
     Validates the inputs that construct a KeyedJaggedTensor.
 
-    This function ensures that:
-    - At least one of lengths or offsets is provided
-    - If both are provided, they are consistent with each other
-    - The dimensions of these tensors align with the values tensor
-
     Any invalid input will result in a ValueError being thrown.
     """
-    # TODO: Add validation checks on keys, values, weights
     _validate_lengths_and_offsets(kjt)
+    _validate_keys(kjt)
+    _validate_weights(kjt)
 
 
 def _validate_lengths_and_offsets(kjt: KeyedJaggedTensor) -> None:
+    """
+    Validates the lengths and offsets of a KJT.
+
+    - At least one of lengths or offsets is provided
+    - If both are provided, they are consistent with each other
+    - The dimensions of these tensors align with the values tensor
+    """
     lengths = kjt.lengths_or_none()
     offsets = kjt.offsets_or_none()
     if lengths is None and offsets is None:
@@ -76,3 +79,45 @@ def _validate_offsets(offsets: torch.Tensor, values: torch.Tensor) -> None:
         raise ValueError(
             f"The last element of offsets must equal to the number of values, but got {offsets[-1]} and {values.numel()}"
         )
+
+
+def _validate_keys(kjt: KeyedJaggedTensor) -> None:
+    """
+    Validates KJT keys, assuming the lengths/offsets input are valid.
+
+    - keys must be unique
+    - For non-VBE cases, the size of lengths is divisible by the number of keys
+    """
+    keys = kjt.keys()
+
+    if len(set(keys)) != len(keys):
+        raise ValueError("keys must be unique")
+
+    lengths = kjt.lengths_or_none()
+    offsets = kjt.offsets_or_none()
+    if lengths is not None:
+        lengths_size = lengths.numel()
+    else:
+        assert offsets is not None
+        lengths_size = offsets.numel() - 1
+
+    if len(keys) == 0 and lengths_size > 0:
+        raise ValueError("keys is empty but lengths or offsets is not")
+    elif len(keys) > 0:
+        # TODO: Validate KJT for VBE cases
+        if not kjt.variable_stride_per_key():
+            if lengths_size % len(keys) != 0:
+                raise ValueError(
+                    f"lengths size must be divisible by keys size, but got {lengths_size} and {len(keys)}"
+                )
+
+
+def _validate_weights(kjt: KeyedJaggedTensor) -> None:
+    """
+    Validates if the KJT weights has the same size as values.
+    """
+    weights = kjt.weights_or_none()
+    if weights is not None and weights.numel() != kjt.values().numel():
+        raise ValueError(
+            f"weights size must equal to values size, but got {weights.numel()} and {kjt.values().numel()}"
+        )

torchrec/sparse/tests/test_jagged_tensor_validator.py

@@ -122,20 +122,71 @@ class TestJaggedTensorValidator(unittest.TestCase):
         ),
     ]
 
-    @parameterized.expand(INVALID_LENGTHS_OFFSETS_CASES)
+    INVALID_KEYS_CASES = [
+        param(
+            expected_error_msg="keys must be unique",
+            keys=["f1", "f1"],
+            values=torch.tensor([1, 2, 3, 4, 5]),
+            lengths=torch.tensor([1, 2, 0, 2]),
+            offsets=torch.tensor([0, 1, 3, 3, 5]),
+        ),
+        param(
+            expected_error_msg="keys is empty but lengths or offsets is not",
+            keys=[],
+            values=torch.tensor([1, 2, 3, 4, 5]),
+            lengths=torch.tensor([1, 2, 0, 2]),
+            offsets=torch.tensor([0, 1, 3, 3, 5]),
+        ),
+        param(
+            expected_error_msg="lengths size must be divisible by keys size",
+            keys=["f1", "f2", "f3"],
+            values=torch.tensor([1, 2, 3, 4, 5]),
+            lengths=torch.tensor([1, 2, 0, 2]),
+            offsets=torch.tensor([0, 1, 3, 3, 5]),
+        ),
+    ]
+
+    INVALID_WEIGHTS_CASES = [
+        param(
+            expected_error_msg="weights size must equal to values size",
+            keys=["f1", "f2"],
+            values=torch.tensor([1, 2, 3, 4, 5]),
+            lengths=torch.tensor([1, 2, 0, 2]),
+            offsets=torch.tensor([0, 1, 3, 3, 5]),
+            weights=torch.tensor([0.1, 0.2, 0.3, 0.4]),
+        ),
+        param(
+            expected_error_msg="weights size must equal to values size",
+            keys=["f1", "f2"],
+            values=torch.tensor([1, 2, 3, 4, 5]),
+            lengths=torch.tensor([1, 2, 0, 2]),
+            offsets=torch.tensor([0, 1, 3, 3, 5]),
+            weights=torch.tensor([0.1, 0.2, 0.3, 0.4, 0.5, 0.6]),
+        ),
+    ]
+
+    @parameterized.expand(
+        [
+            *INVALID_LENGTHS_OFFSETS_CASES,
+            *INVALID_KEYS_CASES,
+            *INVALID_WEIGHTS_CASES,
+        ]
+    )
     def test_invalid_keyed_jagged_tensor(
         self,
         expected_error_msg: str,
         keys: List[str],
         values: torch.Tensor,
         lengths: Optional[torch.Tensor],
         offsets: Optional[torch.Tensor],
+        weights: Optional[torch.Tensor] = None,
     ) -> None:
         kjt = KeyedJaggedTensor(
             keys=keys,
             values=values,
             lengths=lengths,
             offsets=offsets,
+            weights=weights,
         )
 
         with self.assertRaises(ValueError) as err:
@@ -181,3 +232,8 @@ def test_valid_kjt_from_offsets(
         )
 
         validate_keyed_jagged_tensor(kjt)
+
+    def test_valid_empty_kjt(self) -> None:
+        kjt = KeyedJaggedTensor.empty()
+
+        validate_keyed_jagged_tensor(kjt)