add validation for weights tensor in sparse feature (pytorch#680)

Summary:
Pull Request resolved: pytorch#680

Validate weights tensor if it's provided.

Reviewed By: zyan0

Differential Revision: D40031144

fbshipit-source-id: ce6bc715c3ab8b06999c59d8c1a70415e819cd38

Author: 842974287

torchrec/inference/include/torchrec/inference/Validation.h

@@ -16,7 +16,11 @@ namespace torchrec {
 // Currently validates:
 //  1. Whether sum(lengths) == size(values)
 //  2. Whether there are negative values in lengths
-bool validateSparseFeatures(at::Tensor& values, at::Tensor& lengths);
+//  3. If weights is present, whether sum(lengths) == size(weights)
+bool validateSparseFeatures(
+    at::Tensor& values,
+    at::Tensor& lengths,
+    c10::optional<at::Tensor> maybeWeights = c10::nullopt);
 
 // Returns whether dense features are valid.
 // Currently validates:

torchrec/inference/src/Validation.cpp

@@ -11,14 +11,20 @@
 
 namespace torchrec {
 
-bool validateSparseFeatures(at::Tensor& values, at::Tensor& lengths) {
+bool validateSparseFeatures(
+    at::Tensor& values,
+    at::Tensor& lengths,
+    c10::optional<at::Tensor> maybeWeights) {
   auto flatLengths = lengths.view(-1);
 
-  // validate sum of lengths equals number of values
+  // validate sum of lengths equals number of values/weights
   auto lengthsTotal = at::sum(flatLengths).item<int>();
   if (lengthsTotal != values.size(0)) {
     return false;
   }
+  if (maybeWeights.has_value() && lengthsTotal != maybeWeights->size(0)) {
+    return false;
+  }
 
   // Validate no negative values in lengths.
   // Use faster path if contiguous.

torchrec/inference/tests/ValidationTest.cpp

@@ -14,21 +14,24 @@
 TEST(ValidationTest, validateSparseFeatures) {
   auto values = at::tensor({1, 2, 3, 4});
   auto lengths = at::tensor({1, 1, 1, 1});
+  auto weights = at::tensor({.1, .2, .3, .4});
 
   // pass 1D
-  EXPECT_TRUE(torchrec::validateSparseFeatures(values, lengths));
+  EXPECT_TRUE(torchrec::validateSparseFeatures(values, lengths, weights));
 
   // pass 2D
   lengths.reshape({2, 2});
-  EXPECT_TRUE(torchrec::validateSparseFeatures(values, lengths));
+  EXPECT_TRUE(torchrec::validateSparseFeatures(values, lengths, weights));
 
   // fail 1D
   auto invalidLengths = at::tensor({1, 2, 1, 1});
-  EXPECT_FALSE(torchrec::validateSparseFeatures(values, invalidLengths));
+  EXPECT_FALSE(
+      torchrec::validateSparseFeatures(values, invalidLengths, weights));
 
   // fail 2D
   invalidLengths.reshape({2, 2});
-  EXPECT_FALSE(torchrec::validateSparseFeatures(values, invalidLengths));
+  EXPECT_FALSE(
+      torchrec::validateSparseFeatures(values, invalidLengths, weights));
 }
 
 TEST(ValidationTest, validateDenseFeatures) {