MCPcopy Create free account
hub / github.com/apache/tvm / cumsum

Function cumsum

python/tvm/topi/gpu/scan.py:679–726  ·  view source on GitHub ↗

Numpy style cumsum op. Return the cumulative sum of the elements along a given axis. Parameters ---------- data : tvm.te.Tensor The input data to the operator. axis : int, optional Axis along which the cumulative sum is computed. The default (None) is to compute

(
    data: tvm.te.Tensor,
    axis: int | None = None,
    dtype: int | None = None,
    exclusive: bool | None = None,
    workspace: tvm.te.Tensor | None = None,
)

Source from the content-addressed store, hash-verified

677
678
679def cumsum(
680 data: tvm.te.Tensor,
681 axis: int | None = None,
682 dtype: int | None = None,
683 exclusive: bool | None = None,
684 workspace: tvm.te.Tensor | None = None,
685) -> tvm.te.Tensor:
686 """Numpy style cumsum op. Return the cumulative sum of the elements along a given axis.
687
688 Parameters
689 ----------
690 data : tvm.te.Tensor
691 The input data to the operator.
692
693 axis : int, optional
694 Axis along which the cumulative sum is computed. The default (None) is to compute
695 the cumsum over the flattened array.
696
697 dtype : string, optional
698 Type of the returned array and of the accumulator in which the elements are summed.
699 If dtype is not specified, it defaults to the dtype of data.
700
701 exclusive : bool, optional
702 If true will return exclusive sum in which the first element is not
703 included. In other terms, if true, the j-th output element would be
704 the sum of the first (j-1) elements. Otherwise, it would be the sum of
705 the first j elements.
706
707 workspace: Optional[tvm.te.Tensor]
708 A buffer to store intermediate results if thrust is enabled. The size of the workspace
709 should be sufficiently large, this can be obtained by overestimation or memory usage
710 profiling. If None, it will fallback to use thrust internal memory allocation.
711
712 Returns
713 -------
714 result : tvm.te.Tensor
715 The result has the same size as data, and the same shape as data if axis is not None.
716 If axis is None, the result is a 1-d array.
717 """
718 return scanop(
719 data=data,
720 binop=tvm.tirx.generic.add,
721 identity_value=0,
722 axis=axis,
723 dtype=dtype,
724 exclusive=exclusive,
725 workspace=workspace,
726 )
727
728
729def cumprod(

Callers

nothing calls this directly

Calls 1

scanopFunction · 0.70

Tested by

no test coverage detected

Used in the wild real call sites across dependent graphs

searching dependent graphs…