deepmd.pd.utils.decomp

deepmd.pd.utils.decomp#

Attributes#

masked_add_

scatter_reduce

Functions#

sec(→ list[int])

Auxiliary function for decomposed functions.

numel(→ int)

Module Contents#

deepmd.pd.utils.decomp.sec(length: int, size: int) → list[int][source]#

Auxiliary function for decomposed functions.

If length is not divisible by size, the last chunk will be smaller.

Parameters:
lengthint

Length to be chunked.

sizeint

Chunk size.

Returns:
list[int]

Chunked output list.

deepmd.pd.utils.decomp.numel(x: paddle.Tensor) → int[source]#
deepmd.pd.utils.decomp.masked_add_[source]#
deepmd.pd.utils.decomp.scatter_reduce[source]#
Contents