Add check for extra values during unflattening (meta-pytorch#3301)

Summary:
Pull Request resolved: meta-pytorch#3301

As discussed in the previous diff, adding input sanity check

Reviewed By: iamzainhuda

Differential Revision: D80642205

fbshipit-source-id: 8f790f1fcad54a6b0e5cb3432bf251470da6d24e

Author: malaybag

torchrec/sparse/jagged_tensor.py

@@ -3067,6 +3067,10 @@ def _kjt_unflatten(
 ) -> KeyedJaggedTensor:
     if len(values) < len(KeyedJaggedTensor._fields):
         values.extend([None] * (len(KeyedJaggedTensor._fields) - len(values)))
+    elif len(values) > len(KeyedJaggedTensor._fields):
+        raise ValueError(
+            f"Too many values provided for KeyedJaggedTensor: {len(values)} vs {len(KeyedJaggedTensor._fields)}"
+        )
     return KeyedJaggedTensor(
         context,
         *values[:-2],

torchrec/sparse/tests/test_keyed_jagged_tensor.py

@@ -172,6 +172,18 @@ def test_pytree_kjt(self) -> None:
         self.assertIsNone(kjt_0.inverse_indices_or_none())
         self.assertIsNone(kjt_1.inverse_indices_or_none())
 
+        kjt_0 = KeyedJaggedTensor(
+            values=values,
+            keys=keys,
+            offsets=offsets,
+            weights=weights,
+        )
+        elems, spec = pytree.tree_flatten(kjt_0)
+
+        # Simulate extra fields
+        with self.assertRaises(ValueError):
+            kjt_1 = pytree.tree_unflatten(elems + elems, spec)
+
     def test_to_dict_vb(self) -> None:
         values = torch.Tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0, 7.0, 8.0])
         weights = torch.Tensor([1.0, 0.5, 1.5, 1.0, 0.5, 1.0, 1.0, 1.5])