Repository: pytorch/functorch (language: python)

Issue:
[vmap] Most used `torch.<x>` operations: 100-150
- [ ] fft	4436 (needs opinfo)
- [x] lt	4421
- [x] log2	4316
- [x] enable_grad	4238
- [x] rand_like	4187
- [x] argsort	3972
- [x] seed	3932
- [x] mv	3547
- [x] ger	3309
- [x] ge	3248
- [x] atan2	3210
- [x] ceil	3202
- [x] ne	3075
- [ ] bincount	3063 (needs opinfo)
- [x] acos	3055
- [x] rsqrt	3031
- [x] svd	3029
- [x] numel	3003
- [x] log1p	2840
- [x] unbind	2808
- [x] le	2714
- [x] isinf	2707
- [x] cross	2646
- [x] set_default_dtype	2536
- [x] argmin	2535
- [x] ~~sparse_coo_tensor	2489~~
- [x] log10	2304
- [x] kthvalue	2192
- [x] set_rng_state	2158
- [x] get_rng_state	1996
- [x] get_default_dtype	1879
- [x] det	1868
- [x] qr	1864
- [ ] histc	1852 (needs opinfo)
- [x] symeig	1832
- [x] trace	1801
- [x] median	1795
- [x] addcmul	1751
- [x] remainder	1717
- [x] baddbmm	1693
- [x] lgamma	1665
- [ ] repeat_interleave	1598 (actionable)
- [x] fmod	1576
- [x] reciprocal	1575
- [x] tan	1560
- [x] initial_seed	1532
- [ ] take	1529 (actionable)
- [ ] stft	1487 (needs OpInfo)
- [x] get_num_threads	1477
- [x] real	1459


Discussion / hints:

Fix patch:
diff --git a/functorch/csrc/BatchRulesModules.cpp b/functorch/csrc/BatchRulesModules.cpp
index 2e198095a..aa733fa1c 100644
--- a/functorch/csrc/BatchRulesModules.cpp
+++ b/functorch/csrc/BatchRulesModules.cpp
@@ -204,7 +204,7 @@ std::tuple<Tensor,Tensor> cudnn_convolution_backward_plumbing(const Tensor & sel
 /**
  * grid sample batch rule breaks down into 3 cases:
  *   case 1 (input is batched, grid is not):
- *     batch input along first dimension, unpack along first dimension 
+ *     batch input along first dimension, unpack along first dimension
  *     2d:
  *       input: N(BC)H_{in}W_{in}, grid: NH_{out}W_{out}2
  *       output: N(BC)H_{out}W_{out}
@@ -255,6 +255,45 @@ grid_sample_batch_rule(const Tensor& input, optional<int64_t> input_bdim, const
   return result;
 }
 
+std::tuple<Tensor, optional<int64_t>> cross_batch_rule(
+    const Tensor& self, optional<int64_t> self_bdim,
+    const Tensor& other, optional<int64_t> other_bdim,
+    const optional<int64_t> dim) {
+  auto self_ = moveBatchDimToFront(self, self_bdim);
+  auto other_ = moveBatchDimToFront(other, other_bdim);
+
+  if (other_bdim.has_value() && !self_bdim.has_value()) {
+    self_ = self_.expand_as(other_);
+  }
+  if (self_bdim.has_value() && !other_bdim.has_value()) {
+    other_ = other_.expand_as(self_);
+  }
+  auto new_dim = dim;
+  if (dim.has_value()) {
+    auto t = (self_bdim.has_value()) ? self_ : other_;
+    bool flag = (self_bdim.has_value()) ? true : other_bdim.has_value();
+    new_dim = getPhysicalDim(t, flag, *dim);
+  } else {
+    // if batch size is 3 we have to avoid that bdim is used as cross' dim argument
+    // according to cross API:
+    // > If dim is not given, it defaults to the first dimension found with the size 3
+    // we have to skip batch dim and find another dim with size 3
+    auto bs = (self_bdim.has_value()) ? self_.size(0) : (other_bdim.has_value()) ? other_.size(0) : -1;
+    if (bs == 3) {
+      auto t = (self_bdim.has_value()) ? self_ : other_;
+      int64_t idx = 1;
+      for (auto it = t.sizes().begin() + 1; it < t.sizes().end(); ++it, ++idx) {
+        if (*it == 3) {
+          new_dim = idx;
+          break;
+        }
+      }
+    }
+  }
+  optional<int64_t> out_dim = (self_bdim.has_value() || other_bdim.has_value()) ? 0 : (optional<int64_t>) nullopt;
+  return std::make_tuple(at::cross(self_, other_, new_dim), out_dim);
+}
+
 // TODO: replace with targetable functionalization
 Tensor one_hot_decomposition_hack(const Tensor &self, int64_t num_classes) {
     TORCH_CHECK(self.dtype() == kLong, "one_hot is only applicable to index tensor.");
@@ -360,6 +399,7 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   VMAP_SUPPORT("grid_sampler_2d", GRID_SAMPLE_BATCH_RULE(grid_sampler));
   VMAP_SUPPORT("grid_sampler_3d", GRID_SAMPLE_BATCH_RULE(grid_sampler));
   VMAP_SUPPORT("cudnn_grid_sampler", GRID_SAMPLE_BATCH_RULE(cudnn_grid_sampler));
+  VMAP_SUPPORT("cross", cross_batch_rule);
 
   UNARY_POINTWISE(constant_pad_nd);
   EXISTING_BDIM(reflection_pad1d);


Test patch:
diff --git a/test/common_utils.py b/test/common_utils.py
index 2eedb425a..a9dda00a5 100644
--- a/test/common_utils.py
+++ b/test/common_utils.py
@@ -22,7 +22,6 @@ def loop(op, in_dims, out_dim, batch_size, *batched_args, **kwarg_values):
         idx_kwargs = {}
         flat_args, args_spec = pytree.tree_flatten(batched_args)
         flat_dims, dims_spec = pytree.tree_flatten(in_dims)
-        # print(flat_args)
         assert(args_spec == dims_spec)
         new_args = [a.select(in_dim, idx) if in_dim is not None else a for a, in_dim in zip(flat_args, flat_dims)]
         out = op(*pytree.tree_unflatten(new_args, args_spec), **kwarg_values)
@@ -68,7 +67,7 @@ def add_batch_choices(a):
 
 def get_fallback_and_vmap_exhaustive(op, arg_values, kwarg_values, compute_loop_out=True):
     out_dim = 0
-    batch_size = 3
+    batch_size = 4
     generator = get_exhaustive_batched_inputs(arg_values, kwarg_values, batch_size)
     for batched_args, in_dims, kwarg_values in generator:
         if compute_loop_out:
@@ -94,13 +93,13 @@ def f(x, *args, **kwargs):
                 out[idx] = out[idx] + x.to(out[idx].device)
             return out
 
-        vmap1_dims = tuple([0] + [None] * len(in_dims))
-        vmap2_dims = tuple([None] + list(in_dims))
+        vmap1_dims = tuple([None] + list(in_dims))
+        vmap2_dims = tuple([0] + [None] * len(in_dims))
         if compute_loop_out:
-            loop_out = pytree.tree_map(lambda v: torch.ones(3, *v.shape, device=v.device) + v, loop_out)
+            loop_out = pytree.tree_map(lambda v: torch.ones(2, *v.shape, device=v.device) + v, loop_out)
         else:
             loop_out = None
-        batched_out = vmap(vmap(f, in_dims=vmap1_dims), in_dims=vmap2_dims)(torch.ones(3), *batched_args, **kwarg_values)
+        batched_out = vmap(vmap(f, in_dims=vmap1_dims), in_dims=vmap2_dims)(torch.ones(2), *batched_args, **kwarg_values)
         yield (loop_out, batched_out)
 
 def opinfo_in_dict(opinfo, d):
diff --git a/test/test_ops.py b/test/test_ops.py
index a08c2167c..abf8be1bb 100644
--- a/test/test_ops.py
+++ b/test/test_ops.py
@@ -378,7 +378,6 @@ def test_vmapvjp(self, device, dtype, op):
         xfail('complex'),
         xfail('copysign'),
         xfail('corrcoef'),
-        xfail('cross'),
         xfail('cummax'),
         xfail('cummin'),
         xfail('cumprod'),
diff --git a/test/test_vmap.py b/test/test_vmap.py
index 82b9b91c3..434f15d77 100644
--- a/test/test_vmap.py
+++ b/test/test_vmap.py
@@ -1634,6 +1634,17 @@ def test_comparison_ops(self):
             number = getter([]).item()
             self._test_unary(lambda t: op(t, number), getter, 'cpu', check_propagates_grad=False)
 
+    def test_cross_batch_size_three(self):
+        # Let's test corner case when batch_size is 3 and cross' dim argument is not specified
+        # According to the cross API, dim will be assigned to the first dim with value 3
+        # In this test we ensure that found dim is not batch dim.
+        op = torch.cross
+        test = self._vmap_test
+        B0 = B1 = 3
+        test(op, (torch.rand(B0, 2, 3), torch.rand(B0, 2, 3)))
+        test(vmap(op, in_dims=(0, None)), (torch.rand(B0, B1, 2, 3), torch.rand(B0, B1, 2, 3)),
+             in_dims=(None, 1))
+
     def test_diagonal(self):
         tensor = torch.randn(3, 5, 7, 11, 13)
         test = self._vmap_view_test
@@ -2974,7 +2985,6 @@ def test_vmap_exhaustive(self, device, dtype, op):
         xfail('cdist'),
         xfail('complex'),
         xfail('copysign'),
-        xfail('cross'),
         xfail('diag_embed'),
         xfail('dsplit'),
         xfail('eig'),