def init(*args, **kwargs):
_basics.init(*args, **kwargs)
# Call set up again to make sure the basics is in sync
_setup_process_sets(_basics)
def shutdown(*args, **kwargs):
mpi_lib.horovod_torch_reset()
return _basics.shutdown(*args, **kwargs)
# import reduction op values
Average = _basics.Average
Sum = _basics.Sum
Adasum = _basics.Adasum
is_homogeneous = _basics.is_homogeneous
handle_average_backwards_compatibility = get_average_backwards_compatibility_fun(
_basics)
_setup_process_sets(_basics)
# Schema: handle -> input, output
# We keep input in order to make sure it does not get garbage collected
# before the operation is finished.
_handle_map = {}
def _check_function(function_factory, tensor):
function = function_factory(tensor)
if not hasattr(mpi_lib, function):
raise ValueError('Tensor type %s is not supported.' % tensor.type())
if not tensor.is_contiguous():
raise ValueError('Tensor is required to be contiguous.')
return function
def init(*args, **kwargs):
global _handle_map
_handle_map = {}
_basics.init(*args, **kwargs)
# Call set up again to make sure the basics is in sync
_setup_process_sets(_basics)
