Slab#

class nvmath.distributed.distribution.Slab(partition_dim: int, ndim: int | None = None)[source]#

Slab distribution

Data is partitioned across processes on a single axis, such that:

  • The shape of the slab on the first s_p % P processes is (s_0, …, s_p // P + 1, …, s_{n-1})

  • The shape of the slab on the remaining processes is (s_0, …, s_p // P, …, s_{n-1})

  • Process 0 owns the first slab according to the global index order, process 1 owns the second slab and so on.

where:

  • s_i is the size of dimension i of the global array

  • p is the partition dimension

  • n is the number of dimensions of the array

  • P is the number of processes

Methods

__init__(partition_dim: int, ndim: int | None = None)[source]#
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. Must be compatible with the dimensionality of the source distribution. This may be required if the source distribution doesn’t have associated dimensionality.

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

Raises:

ConvertDistributionError – if the conversion is not possible.

Attributes

X: Slab = <nvmath.distributed.distribution.Slab object>#

Slab distribution on axis 0.

Y: Slab = <nvmath.distributed.distribution.Slab object>#

Slab distribution on axis 1.

Z: Slab = <nvmath.distributed.distribution.Slab object>#

Slab distribution on axis 2.

name#
ndim#
partition_dim#

Slab partition dimension