# Buggy code
def allgather(tensor, name=None, priority=0):
    """
    A function that concatenates the input tensor with the same input tensor on
    all other Horovod processes. The input tensor is not modified.

    The concatenation is done on the first dimension, so the input tensors on
    the different processes must have the same rank and shape, except for the
    first dimension, which is allowed to be different.

    This acts as a thin wrapper around an autograd function.  If your input
    tensor requires gradients, then callings this function will allow gradients
    to be computed and backpropagated.

    Arguments:
        tensor: A tensor to allgather.
        name: A name of the allgather operation.
        priority: The priority of this operation. Higher priority operations
                  are likely to be executed before other operations.

    Returns:
        A tensor of the same type as `tensor`, concatenated on dimension zero
        across all processes. The shape is identical to the input shape, except
        for the first dimension, which may be greater and is the sum of all
        first dimensions of the tensors in different Horovod processes.
    """
    assert isinstance(tensor, mx.nd.NDArray)
    output = mx.nd.zeros(shape=tensor.shape, ctx=tensor.context, dtype=tensor.dtype)
    c_in = tensor.handle
    c_out = output.handle
    if isinstance(name, string_types):
        check_call(
            MPI_MXNET_LIB_CTYPES.horovod_mxnet_allgather_async(
                c_in, c_out, c_str(name), ctypes.c_int(priority)
            )
        )
    else:
        check_call(
            MPI_MXNET_LIB_CTYPES.horovod_mxnet_allgather_async(
                c_in, c_out, name, ctypes.c_int(priority)
            )
        )
    return output


# Traceback (mxnet.base.MXNetError)
(before) a.shape = (10, 100, 100)
(before) a[0]=
[[0.6686509  0.17409194 0.3850025  ... 0.43011498 0.0661214  0.2502998 ]
[0.7005292  0.19000232 0.6673837  ... 0.27718288 0.16084558 0.223108  ]
[0.96042585 0.81086403 0.54152083 ... 0.5650488  0.5196334  0.6767488 ]
...
[0.96879214 0.9387428  0.04036242 ... 0.13176239 0.3436321  0.47154343]
[0.8069018  0.91234195 0.01141495 ... 0.35816687 0.57390726 0.68393874]
[0.72049534 0.67948174 0.44702923 ... 0.87448525 0.63809574 0.7006303 ]]
<NDArray 100x100 @gpu(0)>
(after) a.shape = (10, 100, 100)
(before) a.shape = (10, 100, 100)
(before) a[0]=
[[0.7685592  0.10232276 0.8685353  ... 0.93769354 0.62144864 0.21535844]
[0.85973674 0.3420865  0.6202223  ... 0.5464046  0.41442537 0.32170743]
[0.11786121 0.23281038 0.95843846 ... 0.17739207 0.5901362  0.28355032]
...
[0.70749086 0.61171615 0.37854642 ... 0.3485002  0.29636437 0.7359518 ]
[0.4345038  0.90834665 0.5242443  ... 0.0793817  0.40161872 0.6579807 ]
[0.8531954  0.18177992 0.7053579  ... 0.5257004  0.24457276 0.74836564]]
<NDArray 100x100 @gpu(1)>
(after) a.shape = (10, 100, 100)
Traceback (most recent call last):
File "all_gather_bug.py", line 29, in <module>
print("(after) a[0]={}".format(a[0]))
File "/opt/mxnet/python/mxnet/ndarray/ndarray.py", line 194, in __repr__
return '\n%s\n<%s %s @%s>' % (str(self.asnumpy()),
File "/opt/mxnet/python/mxnet/ndarray/ndarray.py", line 1996, in asnumpy
ctypes.c_size_t(data.size)))
File "/opt/mxnet/python/mxnet/base.py", line 252, in check_call
raise MXNetError(py_str(_LIB.MXGetLastError()))
mxnet.base.MXNetError: MPI_Allgatherv failed, see MPI output for details.
Traceback (most recent call last):
File "all_gather_bug.py", line 29, in <module>
print("(after) a[0]={}".format(a[0]))
File "/opt/mxnet/python/mxnet/ndarray/ndarray.py", line 194, in __repr__
return '\n%s\n<%s %s @%s>' % (str(self.asnumpy()),
File "/opt/mxnet/python/mxnet/ndarray/ndarray.py", line 1996, in asnumpy
ctypes.c_size_t(data.size)))
File "/opt/mxnet/python/mxnet/base.py", line 252, in check_call
raise MXNetError(py_str(_LIB.MXGetLastError()))
mxnet.base.MXNetError: MPI_Allgatherv failed, see MPI output for details.
--------------------------------------------------------------------------
Primary job  terminated normally, but 1 process returned
a non-zero exit code. Per user-direction, the job has been aborted.
--------------------------------------------------------------------------
--------------------------------------------------------------------------
mpiexec detected that one or more processes exited with non-zero status, thus causing
the job to be terminated. The first process to do so was:

Process name: [[59002,1],1]
Exit code:    1

# Fixed code
def allgather(tensor, name=None, priority=0):
    """
    A function that concatenates the input tensor with the same input tensor on
    all other Horovod processes. The input tensor is not modified.

    The concatenation is done on the first dimension, so the input tensors on
    the different processes must have the same rank and shape, except for the
    first dimension, which is allowed to be different.

    This acts as a thin wrapper around an autograd function.  If your input
    tensor requires gradients, then callings this function will allow gradients
    to be computed and backpropagated.

    Arguments:
        tensor: A tensor to allgather.
        name: A name of the allgather operation.
        priority: The priority of this operation. Higher priority operations
                  are likely to be executed before other operations.

    Returns:
        A tensor of the same type as `tensor`, concatenated on dimension zero
        across all processes. The shape is identical to the input shape, except
        for the first dimension, which may be greater and is the sum of all
        first dimensions of the tensors in different Horovod processes.
    """
    assert isinstance(tensor, mx.nd.NDArray)
    # Size of output is unknown, create output array that
    # will be resized during Horovod operation
    output = mx.nd.empty(shape=[1], ctx=tensor.context, dtype=tensor.dtype)
    c_in = tensor.handle
    c_out = output.handle
    if isinstance(name, string_types):
        check_call(
            MPI_MXNET_LIB_CTYPES.horovod_mxnet_allgather_async(
                c_in, c_out, c_str(name), ctypes.c_int(priority)
            )
        )
    else:
        check_call(
            MPI_MXNET_LIB_CTYPES.horovod_mxnet_allgather_async(
                c_in, c_out, name, ctypes.c_int(priority)
            )
        )

    # Need to block here so changes to output tensor are visible
    output.wait_to_read()
    return output