Spaces:
Running
Running
| # Copyright (c) ONNX Project Contributors | |
| # SPDX-License-Identifier: Apache-2.0 | |
| import numpy as np | |
| from onnx.reference.op_run import OpRun | |
| class CumSum(OpRun): | |
| def _run(self, x, *axis, exclusive=None, reverse=None): # type: ignore | |
| axis = None if not axis else axis[0] # type: ignore | |
| if axis is None: # type: ignore | |
| if reverse or exclusive: | |
| raise NotImplementedError("reverse=1 or exclusive=1 not implemented") | |
| return (np.cumsum(x),) | |
| if not isinstance(axis, (np.int32, np.int64)): | |
| if len(axis.shape) > 1 or (len(axis.shape) > 0 and axis.shape[0] != 1): # type: ignore | |
| raise RuntimeError( | |
| f"axis must be an array of one number not {axis} (shape {axis.shape})." # type: ignore | |
| ) | |
| if len(axis.shape) > 0: # type: ignore | |
| axis = axis[0] | |
| if reverse: | |
| rev_indices = [slice(0, s) for s in x.shape] | |
| rev_indices[axis] = slice(None, None, -1) # type: ignore | |
| x = x[tuple(rev_indices)] | |
| if exclusive: | |
| indices_c = [slice(0, s) for s in x.shape] | |
| indices_d = [slice(0, s) for s in x.shape] | |
| indices_c[axis] = slice(0, -1) # type: ignore | |
| indices_d[axis] = slice(1, x.shape[axis]) # type: ignore | |
| res = np.zeros(x.shape, dtype=x.dtype) | |
| np.cumsum(x[tuple(indices_c)], axis=axis, out=res[tuple(indices_d)]) # type: ignore | |
| else: | |
| res = np.cumsum(x, axis=axis) # type: ignore | |
| if reverse: | |
| res = res[tuple(rev_indices)] | |
| return (res,) | |