remove pt2 compliant xfails for jagged ops

Summary: letting CI tell me what tests to run to fix these ops for pt2

Differential Revision: D85630006

Author: bdhirsh

fbgemm_gpu/src/jagged_tensor_ops/jagged_tensor_ops.cu

@@ -9,10 +9,6 @@
 #include "common.cuh"
 
 FBGEMM_OP_DISPATCH(CUDA, "dense_to_jagged", fbgemm_gpu::dense_to_jagged);
-FBGEMM_OP_DISPATCH(
-    CUDA,
-    "jagged_to_padded_dense",
-    fbgemm_gpu::jagged_to_padded_dense);
 FBGEMM_OP_DISPATCH(
     CUDA,
     "jagged_dense_elementwise_add",

fbgemm_gpu/src/jagged_tensor_ops/jagged_tensor_ops_autograd.cpp

@@ -785,14 +785,30 @@ class JaggedSliceOp : public torch::autograd::Function<JaggedSliceOp> {
 } // namespace
 
 ///@ingroup jagged-tensor-ops-cpu
-Tensor jagged_to_padded_dense(
+Tensor jagged_to_padded_dense_forward_autograd(
     const Tensor& values,
     const std::vector<Tensor>& offsets,
     const c10::SymIntArrayRef max_lengths,
     const double padding_value) {
   return JaggedToPaddedDenseOp::apply(
       values, offsets, max_lengths, padding_value)[0];
 }
+Tensor jagged_to_padded_dense(
+    const Tensor& values,
+    const std::vector<Tensor>& offsets,
+    const c10::SymIntArrayRef max_lengths,
+    const double padding_value) {
+  static auto op =
+      c10::Dispatcher::singleton()
+          .findSchemaOrThrow("fbgemm::jagged_to_padded_dense_forward", "")
+          .typed<at::Tensor(
+              const Tensor& values,
+              const std::vector<Tensor>& offsets,
+              at::ArrayRef<at::SymInt> max_lengths,
+              const double padding_value)>();
+  Tensor output = op.call(values, offsets, max_lengths, padding_value);
+  return output;
+}
 
 ///@ingroup jagged-tensor-ops-cpu
 /// Output = x + y where x is jagged, y and output are dense
@@ -973,8 +989,16 @@ TORCH_LIBRARY_IMPL(fbgemm, Autograd, m) {
   m.impl("jagged_jagged_bmm", TORCH_FN(fbgemm_gpu::jagged_jagged_bmm));
   m.impl("jagged_dense_bmm", TORCH_FN(fbgemm_gpu::jagged_dense_bmm));
   m.impl("jagged_slice", TORCH_FN(fbgemm_gpu::jagged_slice));
+  m.impl(
+      "jagged_to_padded_dense_forward",
+      TORCH_FN(fbgemm_gpu::jagged_to_padded_dense_forward_autograd));
 }
 
+// These ops are all custom autograd::Functions, which we are registering
+// to the Autograd key above.
+// The only reason that we *also* register to them to the CompositeImplicit
+// key is so that they will decompose by default when using torch.export
+// (even under inference_mode).
 TORCH_LIBRARY_IMPL(fbgemm, CompositeImplicitAutograd, m) {
   m.impl("jagged_index_select", TORCH_FN(fbgemm_gpu::jagged_index_select_2d));
   m.impl("dense_to_jagged", TORCH_FN(fbgemm_gpu::dense_to_jagged));

fbgemm_gpu/src/jagged_tensor_ops/jagged_tensor_ops_cpu.cpp

@@ -1821,10 +1821,9 @@ TORCH_LIBRARY_IMPL(fbgemm, CPU, m) {
   DISPATCH_TO_CPU("dense_to_jagged", fbgemm_gpu::dense_to_jagged);
   DISPATCH_TO_CPU(
       "dense_to_jagged_forward", fbgemm_gpu::dense_to_jagged_forward);
-  DISPATCH_TO_CPU("jagged_to_padded_dense", fbgemm_gpu::jagged_to_padded_dense);
   DISPATCH_TO_CPU(
       "jagged_to_padded_dense_forward",
-      fbgemm_gpu::jagged_to_padded_dense_forward);
+      fbgemm_gpu::jagged_to_padded_dense_forward_cpu);
   DISPATCH_TO_CPU(
       "jagged_to_padded_dense_backward",
       fbgemm_gpu::jagged_to_padded_dense_backward);

fbgemm_gpu/src/jagged_tensor_ops/jagged_tensor_ops_meta.cpp

@@ -53,18 +53,21 @@ Tensor jagged_to_padded_dense_meta(
 
 Tensor jagged_to_padded_dense_backward_meta(
     const at::Tensor& grad_output,
-    const std::vector<Tensor>& /*offsets*/,
+    const std::vector<Tensor>& offsets,
     at::SymInt total_L) {
   const auto& grad_padded_values = grad_output;
 
-  at::SymInt D = grad_padded_values.sym_size(-1);
+  const bool D_folded = grad_padded_values.dim() == offsets.size() + 1;
+  const auto& grad_padded_values_view =
+      D_folded ? grad_padded_values.unsqueeze(-1) : grad_padded_values;
+  at::SymInt D = grad_padded_values_view.sym_size(-1);
   // Initialize with zeros so output will be zero for the portion truncated
   // in forward.
   auto grad_values =
       at::zeros_symint({std::move(total_L), D}, grad_padded_values.options());
 
   TORCH_CHECK(grad_values.is_meta());
-  return grad_values;
+  return D_folded ? grad_values.squeeze(-1) : grad_values;
 }
 
 Tensor jagged_dense_dense_elementwise_add_jagged_output_forward_meta(

fbgemm_gpu/test/jagged/common.py

@@ -43,15 +43,7 @@
 # Please avoid putting tests here, you should put operator-specific
 # skips and failures in deeplearning/fbgemm/fbgemm_gpu/test/failures_dict.json
 # pyre-ignore[24]: Generic type `Callable` expects 2 type parameters.
-additional_decorators: dict[str, list[Callable]] = {
-    "test_pt2_compliant_tag_fbgemm_jagged_dense_elementwise_add": [
-        # This operator has been grandfathered in. We need to fix this test failure.
-        unittest.expectedFailure,
-    ],
-    "test_pt2_compliant_tag_fbgemm_jagged_to_padded_dense": [
-        unittest.expectedFailure,
-    ],
-}
+additional_decorators: dict[str, list[Callable]] = {}
 
 
 def lengths_to_segment_ids(lengths: torch.Tensor) -> torch.Tensor:

fbgemm_gpu/test/jagged/jagged_to_padded_dense_test.py

@@ -113,6 +113,50 @@ def test_jagged_to_padded_dense(
             rtol=1e-3,
         )
 
+        class Mod(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+
+            def forward(self, a, b, c, d):
+                return torch.ops.fbgemm.jagged_to_padded_dense(a, b, c, d)
+
+        with torch.inference_mode():
+            gm = torch.export.export(
+                Mod(),
+                (
+                    x_values.float().requires_grad_(True),
+                    x_offsets,
+                    max_lengths.astype(int).tolist(),
+                    padding_value,
+                ),
+            ).run_decompositions()
+        num_fw_ops = len(
+            [
+                x
+                for x in gm.graph.nodes
+                if x.target is torch.ops.fbgemm.jagged_to_padded_dense_forward.default
+            ]
+        )
+        num_composite_ops = len(
+            [
+                x
+                for x in gm.graph.nodes
+                if x.target is torch.ops.fbgemm.jagged_to_padded_dense.default
+            ]
+        )
+        self.assertEqual(num_fw_ops, 1)
+        self.assertEqual(num_composite_ops, 0)
+
+        torch.library.opcheck(
+            torch.ops.fbgemm.jagged_to_padded_dense,
+            (
+                x_values.float().requires_grad_(True),
+                x_offsets,
+                max_lengths,
+                padding_value,
+            ),
+        )
+
     @given(
         num_jagged_dim=st.integers(1, 5),
         outer_dense_size=st.integers(0, 5),