Update op db and fix failing tests

Author: Chillee

test/functorch_lagging_op_db.py

@@ -13,6 +13,8 @@
 # To achieve this, we keep our OpInfo library behind that of Pytorch's and
 # we periodically update our OpInfo library by regenerating this file
 _functorch_lagging_meta = {
+    ('H', ''),
+    ('T', ''),
     ('__getitem__', ''),
     ('__radd__', ''),
     ('__rand__', ''),
@@ -24,6 +26,10 @@
     ('__rpow__', ''),
     ('__rsub__', ''),
     ('__rxor__', ''),
+    ('_masked.amax', ''),
+    ('_masked.amin', ''),
+    ('_masked.prod', ''),
+    ('_masked.sum', ''),
     ('abs', ''),
     ('acos', ''),
     ('acosh', ''),
@@ -50,19 +56,26 @@
     ('atan2', ''),
     ('atanh', ''),
     ('baddbmm', ''),
+    ('bfloat16', ''),
+    ('bincount', ''),
     ('bitwise_and', ''),
     ('bitwise_left_shift', ''),
     ('bitwise_not', ''),
     ('bitwise_right_shift', ''),
     ('block_diag', ''),
     ('bmm', ''),
+    ('bool', ''),
     ('broadcast_tensors', ''),
     ('broadcast_to', ''),
+    ('bucketize', ''),
+    ('byte', ''),
     ('cat', ''),
     ('cdist', ''),
     ('ceil', ''),
+    ('char', ''),
     ('cholesky', ''),
     ('cholesky_inverse', ''),
+    ('cholesky_solve', ''),
     ('chunk', ''),
     ('clamp', ''),
     ('clamp', 'scalar'),
@@ -94,10 +107,12 @@
     ('div', 'no_rounding_mode'),
     ('div', 'trunc_rounding'),
     ('dot', ''),
+    ('double', ''),
     ('dsplit', ''),
     ('dstack', ''),
     ('eig', ''),
     ('einsum', ''),
+    ('empty_like', ''),
     ('eq', ''),
     ('erf', ''),
     ('erfc', ''),
@@ -108,19 +123,28 @@
     ('expand_as', ''),
     ('expm1', ''),
     ('fft.fft', ''),
+    ('fft.fft2', ''),
     ('fft.fftn', ''),
     ('fft.hfft', ''),
+    ('fft.hfft2', ''),
+    ('fft.hfftn', ''),
     ('fft.ifft', ''),
+    ('fft.ifft2', ''),
     ('fft.ifftn', ''),
     ('fft.ihfft', ''),
+    ('fft.ihfft2', ''),
+    ('fft.ihfftn', ''),
     ('fft.irfft', ''),
+    ('fft.irfft2', ''),
     ('fft.irfftn', ''),
     ('fft.rfft', ''),
+    ('fft.rfft2', ''),
     ('fft.rfftn', ''),
     ('fill_', ''),
     ('flip', ''),
     ('fliplr', ''),
     ('flipud', ''),
+    ('float', ''),
     ('float_power', ''),
     ('floor', ''),
     ('floor_divide', ''),
@@ -130,12 +154,15 @@
     ('fmod', 'autodiffed'),
     ('frac', ''),
     ('frexp', ''),
+    ('full_like', ''),
     ('gather', ''),
     ('ge', ''),
     ('geqrf', ''),
     ('gradient', ''),
     ('gt', ''),
+    ('half', ''),
     ('histogram', ''),
+    ('histogramdd', ''),
     ('hsplit', ''),
     ('hstack', ''),
     ('hypot', ''),
@@ -151,8 +178,15 @@
     ('index_put', ''),
     ('index_select', ''),
     ('inner', ''),
+    ('int', ''),
     ('inverse', ''),
+    ('isfinite', ''),
     ('isin', ''),
+    ('isinf', ''),
+    ('isnan', ''),
+    ('isneginf', ''),
+    ('isposinf', ''),
+    ('isreal', ''),
     ('kron', ''),
     ('kthvalue', ''),
     ('le', ''),
@@ -171,6 +205,7 @@
     ('linalg.inv', ''),
     ('linalg.inv_ex', ''),
     ('linalg.lstsq', ''),
+    ('linalg.lstsq', 'grad_oriented'),
     ('linalg.matrix_norm', ''),
     ('linalg.matrix_power', ''),
     ('linalg.matrix_rank', ''),
@@ -179,6 +214,7 @@
     ('linalg.norm', ''),
     ('linalg.pinv', ''),
     ('linalg.pinv', 'hermitian'),
+    ('linalg.pinv', 'singular'),
     ('linalg.qr', ''),
     ('linalg.slogdet', ''),
     ('linalg.solve', ''),
@@ -199,10 +235,13 @@
     ('logical_not', ''),
     ('logit', ''),
     ('logsumexp', ''),
+    ('long', ''),
     ('lt', ''),
     ('lu', ''),
     ('lu_solve', ''),
     ('lu_unpack', ''),
+    ('mH', ''),
+    ('mT', ''),
     ('masked_fill', ''),
     ('masked_scatter', ''),
     ('masked_select', ''),
@@ -238,14 +277,19 @@
     ('ne', ''),
     ('neg', ''),
     ('nextafter', ''),
+    ('nn.functional.adaptive_avg_pool1d', ''),
     ('nn.functional.adaptive_avg_pool2d', ''),
+    ('nn.functional.adaptive_avg_pool3d', ''),
+    ('nn.functional.avg_pool1d', ''),
     ('nn.functional.avg_pool2d', ''),
+    ('nn.functional.avg_pool3d', ''),
     ('nn.functional.batch_norm', ''),
     ('nn.functional.batch_norm', 'without_cudnn'),
     ('nn.functional.conv2d', ''),
     ('nn.functional.conv_transpose2d', ''),
     ('nn.functional.cosine_similarity', ''),
     ('nn.functional.dropout', ''),
+    ('nn.functional.embedding', ''),
     ('nn.functional.gelu', ''),
     ('nn.functional.grid_sample', ''),
     ('nn.functional.hardshrink', ''),
@@ -270,6 +314,9 @@
     ('nn.functional.pad', 'constant'),
     ('nn.functional.pad', 'reflect'),
     ('nn.functional.pad', 'replicate'),
+    ('nn.functional.pairwise_distance', ''),
+    ('nn.functional.pixel_shuffle', ''),
+    ('nn.functional.pixel_unshuffle', ''),
     ('nn.functional.relu', ''),
     ('nn.functional.relu6', ''),
     ('nn.functional.softplus', ''),
@@ -278,6 +325,7 @@
     ('norm', 'fro'),
     ('norm', 'inf'),
     ('norm', 'nuc'),
+    ('ones_like', ''),
     ('ormqr', ''),
     ('outer', ''),
     ('permute', ''),
@@ -295,6 +343,7 @@
     ('qr', ''),
     ('quantile', ''),
     ('rad2deg', ''),
+    ('randn_like', ''),
     ('ravel', ''),
     ('real', ''),
     ('reciprocal', ''),
@@ -319,6 +368,7 @@
     ('scatter_add', ''),
     ('select', ''),
     ('sgn', ''),
+    ('short', ''),
     ('sigmoid', ''),
     ('sign', ''),
     ('signbit', ''),
@@ -386,6 +436,7 @@
     ('where', ''),
     ('xlogy', ''),
     ('zero_', ''),
+    ('zeros_like', ''),
 }
 
 

test/test_ops.py

@@ -304,9 +304,7 @@ def vjp_of_vjp(*args_and_cotangents):
             for loop_out, batched_out in \
                     get_fallback_and_vmap_exhaustive(vjp_of_vjp, args_and_cotangents, {}):
                 self.assertEqual(loop_out, batched_out, atol=1e-4, rtol=1e-4)
-
-    @ops(functorch_lagging_op_db + additional_op_db, allowed_dtypes=(torch.float,))
-    @skipOps('TestOperators', 'test_vmapvjp', vjp_fail.union({
+    vmapvjp_fail = vjp_fail.union({
         # All of the following are bugs and need to be fixed
         xfail('clamp', ''),
         xfail('diag_embed'),
@@ -360,7 +358,16 @@ def vjp_of_vjp(*args_and_cotangents):
         xfail('nanmean'),
         xfail('block_diag'),
         xfail('nn.functional.dropout'),
-    }))
+        xfail('double'),
+        xfail('fft.fft2'),
+        xfail('fft.ifft2'),
+        xfail('fft.ihfft2'),
+        xfail('fft.ihfftn'),
+        xfail('fft.rfft2'),
+        xfail('_masked.prod'), # calls aten::item
+    })
+    @ops(functorch_lagging_op_db + additional_op_db, allowed_dtypes=(torch.float,))
+    @skipOps('TestOperators', 'test_vmapvjp', vmapvjp_fail)
     def test_vmapvjp(self, device, dtype, op):
         # These are too annoying to put into the list above
         if op.name in {'nn.functional.linear', 'nn.functional.conv2d'}:
@@ -382,7 +389,7 @@ def test_vmapvjp(self, device, dtype, op):
                 self.assertEqual(loop_out, batched_out, atol=1e-4, rtol=1e-4)
 
     @ops(functorch_lagging_op_db + additional_op_db, allowed_dtypes=(torch.float,))
-    @skipOps('TestOperators', 'test_vmapvjp_has_batch_rule', {
+    @skipOps('TestOperators', 'test_vmapvjp_has_batch_rule', vmapvjp_fail.union({
         xfail('nn.functional.pad', 'constant'),
         xfail('view_as_complex'),
         xfail('__getitem__'),
@@ -493,7 +500,22 @@ def test_vmapvjp(self, device, dtype, op):
         xfail('block_diag'),
         xfail('nn.functional.dropout'),
         xfail('nn.functional.batch_norm'),
-    })
+        xfail('_masked.amax'),
+        xfail('_masked.amin'),
+        xfail('_masked.sum'),
+        xfail('_masked.prod'),
+        xfail('cholesky_solve'),
+        xfail('double'),
+        xfail('fft.fft2'),
+        xfail('fft.ifft2'),
+        xfail('fft.ihfft2'),
+        xfail('fft.ihfftn'),
+        xfail('fft.rfft2'),
+        xfail('nn.functional.adaptive_avg_pool1d'),
+        xfail('nn.functional.adaptive_avg_pool3d'),
+        xfail('nn.functional.avg_pool3d'),
+        xfail('nn.functional.embedding'),
+    }))
     def test_vmapvjp_has_batch_rule(self, device, dtype, op):
         # These are too annoying to put into the list above
         if op.name in {'nn.functional.linear', 'nn.functional.conv2d'}:
@@ -522,6 +544,8 @@ def test():
 
     @ops(functorch_lagging_op_db + additional_op_db, allowed_dtypes=(torch.float,))
     @skipOps('TestOperators', 'test_vjpvmap', vjp_fail.union({
+        # fallback path doesn't work
+        xfail('H'),
         # All of the following are bugs and need to be fixed
         xfail('__getitem__'),
         xfail('clamp', ''),
@@ -543,6 +567,7 @@ def test():
         xfail('lu_unpack'),
         xfail('matrix_exp'),
         xfail('view_as_complex'),
+        xfail('double'),
     }))
     def test_vjpvmap(self, device, dtype, op):
         # NB: there is no vjpvmap_has_batch_rule test because that is almost

test/test_pythonkey.py

@@ -197,6 +197,8 @@ class TestPythonKeyOperatorsOpInfo(TestCase):
     xfail('nn.functional.dropout'),
     xfail('linalg.eigvals'),
     xfail('nn.functional.pad', 'circular'),
+    xfail('empty_like'), # randomness
+    xfail('randn_like'), # randomness
     })
     def test_make_fx_exhaustive(self, device, dtype, op):