Decompose arange.default to arange.start_step (#99739)

The aten op arange.default is not in the core aten IR, and should decompose into the arange.start_step op.
Pull Request resolved: #99739
Approved by: https://github.com/SherlockNoMad

torch/_decomp/__init__.py

@@ -184,6 +184,8 @@ def core_aten_decompositions() -> Dict[OpOverload, Callable]:
             aten.addcmul_,
             aten.addr,
             aten.aminmax,
+            aten.arange.default,
+            aten.arange.start,
             aten.avg_pool2d_backward,
             aten.binary_cross_entropy,
             aten.binary_cross_entropy_backward,

torch/_decomp/decompositions.py

@@ -3356,6 +3356,36 @@ def nansum(self, dim=None, keepdim=False, *, dtype=None):
     return aten.sum(torch.where(torch.isnan(self), 0, self), dim, keepdim, dtype=dtype)
 
 
+@register_decomposition([aten.arange.default, aten.arange.out])
+@out_wrapper()
+def arange_default(
+    end: NumberType,
+    *,
+    dtype: Optional[torch.dtype] = None,
+    layout: torch.layout = torch.strided,
+    device: Optional[torch.device] = None,
+    pin_memory: bool = False,
+):
+    return aten.arange.start_step(
+        0, end, 1, dtype=dtype, layout=layout, device=device, pin_memory=pin_memory
+    )
+
+
+@register_decomposition([aten.arange.start])
+def arange_start(
+    start: NumberType,
+    end: NumberType,
+    *,
+    dtype: Optional[torch.dtype] = None,
+    layout: torch.layout = torch.strided,
+    device: Optional[torch.device] = None,
+    pin_memory: bool = False,
+):
+    return aten.arange.start_step(
+        start, end, 1, dtype=dtype, layout=layout, device=device, pin_memory=pin_memory
+    )
+
+
 def register_inplace(aten_op, outplace_op):
     @register_decomposition(aten_op)
     def inplace_op(*args, **kwargs):

torch/_refs/__init__.py

@@ -4341,7 +4341,7 @@ def empty_like(
     )
 
 
-@register_decomposition(aten.arange)
+@register_decomposition([aten.arange.start_step, aten.arange.start_out])
 @out_wrapper()
 def arange(
     start: NumberType = 0,