The nvmath.distributed.distribution module is experimental and potentially subject to future changes.

Box#

class nvmath.distributed.distribution.Box(
lower: Sequence[int],
upper: Sequence[int],
)[source]#

Box distribution

Attributes

box_shape#
lower#

Box lower coordinates

ndim#
upper#

Box upper coordinates

Methods

copy() → Distribution[source]#

This is a common implementation for those distributions that only require a shallow copy.

shape(
rank: int,
global_shape: Sequence[int] | None = None,
) → tuple[int, ...][source]#

Get the local shape of data on the given rank according to this distribution.

Parameters:

  • rank – the process rank for which to calculate the local shape.

  • global_shape – Global shape of data. Required if the distribution is not bound to a global shape, otherwise not required.

to(cls, /, *, ndim=None, copy=False)[source]#

Convert this distribution object to an equivalent distribution of the given type.

Parameters:

  • cls – the target distribution type.

  • ndim – dimensionality of the target distribution. It must be greater than or equal to the source dimensionality; it can never be smaller. When it is larger, the distribution is promoted by appending trailing size-1 dimensions. Providing ndim is required when the source distribution has no associated dimensionality of its own.

  • copy – Returns a copy if the source and target type are the same.

Raises:

ConvertDistributionError – if the conversion is not possible.