mpi_utils#

Computes the number of un-evenly divided elements per section.

Parameters:

• num_elements (int) –

  The total number of elements to divide.

• num_sections (int, default: size ) –

  The number of sections to divide the elements into. Defaults to the number of MPI ranks.

• strategy (str, default: 'balanced' ) –

  The strategy to use for dividing the elements. Can be one of "balanced" (default) or "greedy". In the "balanced" strategy, the elements are divided as evenly as possible across the sections. In the "greedy" strategy, the elements are divided such that the we get many sections with the maximum number of elements.

Returns:

• section_sizes ( list ) –

  The sizes of each section.

• effective_num_elements ( int ) –

  The effective number of elements after sectioning.

Examples:

>>> get_section_sizes(10, 3, "balanced")
([4, 3, 3], 12)
>>> get_section_sizes(10, 3, "greedy")
([4, 4, 2], 12)

Loads an array from disk and broadcasts it to all ranks.

Parameters:

• path (Path) –

  The path to the file to load.

Returns:

• spmatrix | NDArray | dict –

  The loaded array/s.

Raises:

• FileNotFoundError –

  Occurs on every rank where the file does not exist.

Returns the local slice of a distributed array.

Parameters:

• global_array (NDArray) –

  The global array to slice.

Returns:

• NDArray –

  The local slice of the global array.