Added rot90 batch rule (#138)

* Added rot90 batch rule Description: - Added rot90 batch rule - Enabled associated tests Note: tests seem not be reliable to output bdim. Tests are passing even if batching rule outputs wrong bdim * Replaced manual batching rule with REDUCTION_BOXED_ARGS * Removed commented xfail('rot90')
pytorch · Oct 22, 2021 · a99a86a · a99a86a
1 parent 599c5a9
commit a99a86a
Show file tree

Hide file tree

Showing 3 changed files with 1 addition and 2 deletions.
diff --git a/functorch/csrc/BatchRulesReduceOps.cpp b/functorch/csrc/BatchRulesReduceOps.cpp
@@ -317,6 +317,7 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   REDUCTION_BOXED(var.correction);
   REDUCTION_BOXED(var_mean.correction);
   REDUCTION_BOXED(_log_softmax);
+  REDUCTION_BOXED_ARGS(rot90, 2);
 
   VMAP_SUPPORT("_log_softmax_backward_data", _log_softmax_backward_batch_rule);
   VMAP_SUPPORT("_softmax_backward_data", _softmax_backward_batch_rule);

diff --git a/test/test_ops.py b/test/test_ops.py
@@ -454,7 +454,6 @@ def test_vmapvjp(self, device, dtype, op):
         xfail('quantile'),
         xfail('renorm'),
         xfail('roll'),
-        xfail('rot90'),
         xfail('scatter_add'),
         xfail('solve'),
         xfail('sort'),

diff --git a/test/test_vmap.py b/test/test_vmap.py
@@ -3030,7 +3030,6 @@ def test_vmap_exhaustive(self, device, dtype, op):
         xfail('resolve_conj'),
         xfail('resolve_neg'),
         xfail('roll'),
-        xfail('rot90'),
         xfail('scatter'),
         xfail('scatter_add'),
         xfail('take'),