impl_abstract expand_into_jagged_permute

Differential Revision: D50586828
pytorch · Oct 24, 2023 · 539c47b · 539c47b
1 parent 0f6665c
commit 539c47b
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 0 deletions.
diff --git a/fbgemm_gpu/fbgemm_gpu/sparse_operators.py b/fbgemm_gpu/fbgemm_gpu/sparse_operators.py
@@ -71,3 +71,20 @@ def permute_1D_sparse_data_meta(
         # pyre-fixme
         permuted_weights = weights.new_empty(permuted_indices_size)
     return permuted_lengths, permuted_indices, permuted_weights
+
+
+@torch.library.impl_abstract("fbgemm::expand_into_jagged_permute")
+def expand_into_jagged_permute_meta(
+    permute, input_offsets, output_offsets, output_size
+):
+    torch._check(permute.numel() > 0, lambda: "expected {permute.numel} > 0")
+    torch._check(
+        permute.numel() == input_offsets.numel() - 1,
+        lambda: f"expected {permute.numel()} == {input_offsets.numel()} - 1",
+    )
+    torch._check(
+        permute.numel() == output_offsets.numel() - 1,
+        lambda: f"expected {permute.numel()} == {output_offsets.numel()} - 1",
+    )
+    output_permute = input_offsets.new_empty(output_size)
+    return output_permute
diff --git a/fbgemm_gpu/src/sparse_ops/sparse_ops_cpu.cpp b/fbgemm_gpu/src/sparse_ops/sparse_ops_cpu.cpp
@@ -2699,6 +2699,10 @@ TORCH_LIBRARY_FRAGMENT(fbgemm, m) {
   m.def("invert_permute(Tensor permute) -> Tensor");
   m.def(
       "expand_into_jagged_permute(Tensor permute, Tensor input_offset, Tensor output_offset, int output_size) -> Tensor");
+  m.impl_abstract_pystub(
+      "expand_into_jagged_permute",
+      "fbgemm_gpu.sparse_operators",
+      "//deeplearning/fbgemm/fbgemm_gpu:sparse_operators");
   m.def(
       "block_bucketize_sparse_features(Tensor lengths, Tensor indices, bool bucketize_pos, bool sequence, Tensor block_sizes, int my_size, Tensor? weights=None, Tensor? batch_size_per_feature=None, int max_B= -1) -> (Tensor, Tensor, Tensor?, Tensor?, Tensor?)");
   m.def(