Use torch.cumsum instead of numpy one (#109400)

malfet · pytorchmergebot · commit fb58a72d968e · 2023-09-16T02:52:49.000Z
`s/list(numpy.cumsum(foo))/torch.cumsum(torch.tensor(foo), 0).tolist()/` Test plan: ` python3 ../test/inductor/test_split_cat_fx_passes.py -v` Partially addresses #109387 Pull Request resolved: #109400 Approved by: https://github.com/ezyang
diff --git a/torch/_inductor/fx_passes/split_cat.py b/torch/_inductor/fx_passes/split_cat.py
@@ -3,8 +3,6 @@
 import operator
 from typing import Callable, List, Sequence, Tuple, Union
 
-import numpy
-
 import torch
 from torch._dynamo.utils import counters
 
@@ -454,8 +452,7 @@ def get_simplified_split_ranges(
                 for user_input in user_inputs
                 if isinstance(user_input, tuple)
             }
-
-        cumulative_sizes = [0] + list(numpy.cumsum(split_sections))
+        cumulative_sizes = [0] + torch.cumsum(torch.tensor(split_sections), 0).tolist()
         split_ranges = sorted(
             [(cumulative_sizes[r[0]], cumulative_sizes[r[1] + 1]) for r in ranges]
         )
@@ -578,7 +575,7 @@ def replace_split(
                     for i in range(len(split_ranges))
                 ]
         # Now assign the right getitem to the right input
-        cumulative_sizes = [0] + list(numpy.cumsum(split_sections))
+        cumulative_sizes = [0] + torch.cumsum(torch.tensor(split_sections), 0).tolist()
         new_user_inputs_list = []
         for user_inputs in user_inputs_list:
             new_user_inputs = []
