romtools.linalg.parallel_utils

Helper functions for testing and writing parallel code with the romtools.linalg module.

Functions

assert_axis_is_none_or_within_rank(a, axis)	Simple helper function to assert that the given axis is valid for the given array.
distribute_array_impl(global_array, comm[, ...])	Splts an np.array and distributes to all available MPI processes as evenly as possible.
generate_local_and_global_arrays_from_example_impl(...)	Generates and returns the local and global arrays built from the example tensors in the documentation.
generate_random_local_and_global_arrays_impl(...)	Randomly generates a global array of the specified shape and distributes to all available MPI processes.

romtools.linalg.parallel_utils.assert_axis_is_none_or_within_rank(a, axis)

Simple helper function to assert that the given axis is valid for the given array.

romtools.linalg.parallel_utils.distribute_array_impl(global_array, comm, dist_axis=0)

Splts an np.array and distributes to all available MPI processes as evenly as possible.

For example, distributing an array with 6 rows over 3 processes will result in 2 rows per process. If the array has 7 rows, 2 processors will hold 2 rows, and the third will hold 3 rows.

Inputs:

global_array: The global np.array to be distributed. comm: The MPI communicator dist_axis: The axis along which to split the input array. By default, splits along the first axis (rows).

Returns:

The subset of global_array sent to the current MPI process.

Return type:

local_array

romtools.linalg.parallel_utils.generate_local_and_global_arrays_from_example_impl(rank, slices, example)

Generates and returns the local and global arrays built from the example tensors in the documentation.

Parameters:

example (int)

romtools.linalg.parallel_utils.generate_random_local_and_global_arrays_impl(shape, comm, use_int=False)

Randomly generates a global array of the specified shape and distributes to all available MPI processes.

Returns both the local and global arrays.