Arm backend: Add support for scalar inputs for floor_divide.default

Signed-off-by: Agrima Khare <[email protected]>
Change-Id: Id32fd6df6d7f11801880f7a734d9b43b4fb577dd

Author: agrima1304

backends/arm/_passes/decompose_floor_divide_pass.py

@@ -34,9 +34,15 @@ def get_floor_divide_decomposition(op) -> tuple:
     """
 
     if op in edge_floor_divide_ops:
-        return (exir_ops.edge.aten.div.Tensor_mode,)
+        return (
+            exir_ops.edge.aten.div.Tensor_mode,
+            exir_ops.edge.aten.full_like.default,
+        )
     if op in aten_floor_divide_ops:
-        return (torch.ops.aten.div.Tensor_mode,)
+        return (
+            torch.ops.aten.div.Tensor_mode,
+            torch.ops.aten.full_like.default,
+        )
 
     raise RuntimeError(f"Can't get floor_div decomposition for op {op}")
 
@@ -52,11 +58,16 @@ def call_operator(self, op, args, kwargs, meta):
         if op not in (edge_floor_divide_ops + aten_floor_divide_ops):
             return super().call_operator(op, args, kwargs, meta, updated=False)
 
-        (div_op,) = get_floor_divide_decomposition(op)
+        (div_op, full_op) = get_floor_divide_decomposition(op)
 
         input = args[0]
         other = args[1]
 
+        if isinstance(other, int):
+            other = super().call_operator(
+                full_op, (input, other), {}, meta, updated=False
+            )
+
         div_node = super().call_operator(
             div_op, (input, other), {"rounding_mode": "floor"}, meta, updated=True
         )

backends/arm/test/ops/test_floor_div.py

@@ -36,10 +36,6 @@
         torch.ones(5, 10, 25, 20),
         (-1) * torch.ones(5, 10, 25, 20),
     ),
-    "op_floor_div_rank4_large_rand": lambda: (
-        200 * torch.rand(5, 10, 25, 20),
-        torch.rand(5, 10, 25, 20),
-    ),
     "op_floor_div_rank4_randn_mutltiple_broadcasts": lambda: (
         torch.randn(1, 4, 4, 1),
         torch.randn(1, 1, 4, 4),
@@ -48,6 +44,10 @@
         torch.randn(1, 4, 4, 1),
         2,
     ),
+    "op_floor_div_rank4_large_rand": lambda: (
+        200 * torch.rand(5, 10, 25, 20),
+        torch.rand(5, 10, 25, 20),
+    ),
 }