Update on "add meta for _segment_reduce_backward"

bdhirsh · bdhirsh · commit 9268bb5c249b · 2024-10-08T07:34:12.000-07:00
reland of #124988 [ghstack-poisoned]
diff --git a/test/functorch/test_aotdispatch.py b/test/functorch/test_aotdispatch.py
@@ -6307,12 +6307,6 @@ def fn_(x):
     xfail(
         "nn.functional.nll_loss", ""
     ),  # Cannot call sizes() on tensor with symbolic sizes/strides
-    xfail(
-        "_segment_reduce", "lengths"
-    ),  # aten.segment_reduce.default - couldn't find symbolic meta functio...
-    xfail(
-        "_segment_reduce", "offsets"
-    ),  # aten.segment_reduce.default - couldn't find symbolic meta functio...
     xfail("trace", ""),  # Cannot call sizes() on tensor with symbolic sizes/strides
     xfail(
         "_upsample_bilinear2d_aa"
diff --git a/test/test_meta.py b/test/test_meta.py
@@ -1640,11 +1640,17 @@ def test_segment_reduce_backward(self):
         grad = torch.ones(16, dtype=torch.float)
         output = torch.ones(16, dtype=torch.float)
         data = torch.ones(16, dtype=torch.float)
-        reduce_str='max'
+        reduce_str = 'max'
         lengths = torch.ones(16, dtype=torch.long)
 
         out = torch.ops.aten._segment_reduce_backward(grad, output, data, reduce_str, lengths=lengths)
-        out_meta = torch.ops.aten._segment_reduce_backward(grad.to(device='meta'), output.to(device='meta'), data.to(device='meta'), reduce_str, lengths=lengths.to(device='meta'))
+        out_meta = torch.ops.aten._segment_reduce_backward(
+            grad.to(device='meta'),
+            output.to(device='meta'),
+            data.to(device='meta'),
+            reduce_str,
+            lengths=lengths.to(device='meta'),
+        )
         self.assertEqual(out.shape, out_meta.shape)
         self.assertEqual(out.stride(), out_meta.stride())
         self.assertEqual(out.dtype, out_meta.dtype)
@@ -1654,7 +1660,13 @@ def test_segment_reduce_backward(self):
         grad = torch.ones(16, 2, dtype=torch.float)[:, 1]
         data = torch.ones(16, 2, dtype=torch.float)[:, 1]
         out = torch.ops.aten._segment_reduce_backward(grad, output, data, reduce_str, lengths=lengths)
-        out_meta = torch.ops.aten._segment_reduce_backward(grad.to(device='meta'), output.to(device='meta'), data.to(device='meta'), reduce_str, lengths=lengths.to(device='meta'))
+        out_meta = torch.ops.aten._segment_reduce_backward(
+            grad.to(device='meta'),
+            output.to(device='meta'),
+            data.to(device='meta'),
+            reduce_str,
+            lengths=lengths.to(device='meta'),
+        )
         self.assertEqual(out.shape, out_meta.shape)
         self.assertEqual(out.stride(), out_meta.stride())
         self.assertEqual(out.dtype, out_meta.dtype)
