util.cumsum does not scale

Description

(A clear and concise description of what the feature is.)

• util.cumsum implementation https://github.com/awslabs/gluon-ts/blob/master/src/gluonts/mx/util.py#L326 does not scale under mx.ndarray
• cumsum is 2-5 times slower than nd.cumsum under both mx.sym and mx.ndarray, and even fails for large 4-dim input

Sample test

Code

# import ...
def test_cumsum_util(F, shape, method) -> None:
    x = F.ones(shape)
    F_str = 'mx.sym' if F == mx.sym else 'mx.nd'
    st = time.time()
    try:
        if method=="nd_cumsum":
            F.cumsum(x, axis=-1)
        elif method=="util_cumsum":
            cumsum(F, x)
        print(f"F={F_str}, shape={shape}, method={method}: {time.time()-st} sec")
    except:
        print(f"F={F_str}, shape={shape}, method={method}:  failed")

for shape in [ (1000, 210, 70, 10), (5000, 210, 70, 10)]:
    for F in [mx.sym, mx.nd]:
        for method in ["nd_cumsum", "util_cumsum"]:
            test_cumsum_util(F, shape, method)

Results

F=mx.sym, shape=(1000, 210, 70, 10), method=nd_cumsum: 0.00014519691467285156 sec F=mx.sym, shape=(1000, 210, 70, 10), method=util_cumsum: 0.0002918243408203125 sec F=mx.nd, shape=(1000, 210, 70, 10), method=nd_cumsum: 8.392333984375e-05 sec F=mx.nd, shape=(1000, 210, 70, 10), method=util_cumsum: 0.0002658367156982422 sec F=mx.sym, shape=(5000, 210, 70, 10), method=nd_cumsum: 3.886222839355469e-05 sec F=mx.sym, shape=(5000, 210, 70, 10), method=util_cumsum: 0.00016999244689941406 sec F=mx.nd, shape=(5000, 210, 70, 10), method=nd_cumsum: 3.814697265625e-05 sec F=mx.nd, shape=(5000, 210, 70, 10), method=util_cumsum: failed

References

• https://github.com/awslabs/gluon-ts/blob/master/test/support/test_util.py#L35