add atleast_nd decompositions + tests

Author: Chillee

functorch/csrc/BatchRulesStopDecomposition.cpp renamed to functorch/csrc/BatchRulesDecompositions.cpp

@@ -25,6 +25,12 @@ TORCH_LIBRARY_IMPL(aten, FT_BATCHED_KEY, m) {
   OP_DECOMPOSE(arcsinh);
   OP_DECOMPOSE(arctan);
   OP_DECOMPOSE(arctanh);
+  OP_DECOMPOSE(atleast_1d);
+  OP_DECOMPOSE2(atleast_1d, Sequence);
+  OP_DECOMPOSE(atleast_2d);
+  OP_DECOMPOSE2(atleast_2d, Sequence);
+  OP_DECOMPOSE(atleast_3d);
+  OP_DECOMPOSE2(atleast_3d, Sequence);
   OP_DECOMPOSE(broadcast_tensors);
   OP_DECOMPOSE(broadcast_to);
   OP_DECOMPOSE(clip);

test/functorch_additional_op_db.py

@@ -215,3 +215,27 @@ def sample_inputs_cross_entropy(self, device, dtype, requires_grad, reduction):
                supports_out=True))
 
 
+
+def sample_inputs_atleast_nd(self, device, dtype, requires_grad):
+    inps = []
+    for i in range(5):
+        inps.append(make_tensor(list(range(i)), device=device, dtype=dtype,
+                      requires_grad=requires_grad, low=-1, high=1))
+
+    sample_inputs = []
+    for inp in inps:
+        sample_inputs.append(SampleInput(inp))
+
+    sample_inputs.append(SampleInput(inps))
+    return sample_inputs
+
+for i in range(1, 4):
+    additional_op_db.append(
+        OpInfo(f'atleast_{i}d',
+                aten_name="atleast_{i}d",
+                supports_autograd=True,
+                sample_inputs_func=sample_inputs_atleast_nd,
+                dtypesIfCUDA=floating_types_and(torch.half, torch.bfloat16),
+                supports_out=False))
+
+

test/test_vmap.py

@@ -2951,6 +2951,7 @@ def test_diagonal(self, device):
         x = torch.randn(3, 4, 5, device=device, requires_grad=True)
         self._batched_grad_test(lambda x: x.diagonal(0, -1, -2), (x,))
 
+
     @allowVmapFallbackUsage
     def test_unrelated_output(self, device):
         B0 = 3