Abstract operator for bounds_check_indices (#2495)

Summary:
Pull Request resolved: #2495

Create an abstract impl for fbgemm operator bounds_check_indices

Reviewed By: IvanKobzarev

Differential Revision: D56020789

fbshipit-source-id: e5291bb553614a57d405f4d6c4b901c1a03818b4

fbgemm_gpu/codegen/utils/embedding_bounds_check_host_cpu.cpp

@@ -178,6 +178,9 @@ TORCH_LIBRARY_FRAGMENT(fb, m) {
 TORCH_LIBRARY_FRAGMENT(fbgemm, m) {
   // The (a!) tells PyTorch this is an impure operation and so cannot be CSE'd
   // or DCE'd, etc.
+  m.impl_abstract_pystub(
+      "fbgemm_gpu.sparse_ops",
+      "//deeplearning/fbgemm/fbgemm_gpu:sparse_ops_py");
   m.def(
       "bounds_check_indices(Tensor rows_per_table, Tensor(a!) indices, Tensor(b!) offsets, int bounds_check_mode, Tensor(c!) warning, Tensor(d!)? weights=None, Tensor? B_offsets=None, SymInt max_B=-1) -> ()",
       {PT2_COMPLIANT_TAG});

fbgemm_gpu/fbgemm_gpu/sparse_ops.py

@@ -558,3 +558,19 @@ def keyed_jagged_index_select_dim1_abstract(
         ret.append(weights.new_empty([selected_lengths_sum]))
 
     return ret
+
+
+@impl_abstract("fbgemm::bounds_check_indices")
+def bounds_check_indices_abstract(
+    rows_per_table: torch.Tensor,
+    indices: torch.Tensor,
+    offsets: torch.Tensor,
+    bounds_check_mode_int: int,
+    bounds_check_warning: torch.Tensor,
+    per_sample_weights: Optional[torch.Tensor] = None,
+) -> None:
+    """
+    This meta function is used to fake the bounds checking
+    from the original function `fbgemm::bounds_check_indices`
+    """
+    return