Fix aten_cumsum traced mode (#1605)

aten_cumsum was not traceable when the input is casted because the shape
information was not propagated. Turning the implementation to use pure
tracing instead.

onnxscript/function_libs/torch_lib/ops/core.py

@@ -2416,18 +2416,11 @@ def aten_cumsum(
         cast = self
     else:
         cast = op.Cast(self, to=dtype)
-    return _aten_cumsum_onnx(cast, dim)
-
-
-@torch_op("aten::cumsum", private=True, traceable=True)
-def _aten_cumsum_onnx(
-    self: TRealUnlessInt16OrInt8, dim: Union[INT32, INT64]
-) -> TRealUnlessInt16OrInt8:
-    if IsScalar(self):
+    if len(self.shape) == 0:
         # A scalar
-        result = op.Identity(self)
+        result = op.Identity(cast)
     else:
-        result = op.CumSum(self, dim)
+        result = op.CumSum(cast, dim)
     return result