#!/usr/bin/python
"""
demo python script of rabit
"""
import os
import sys
import numpy as np
# add path to wrapper
sys.path.append(os.path.dirname(__file__) + '/../wrapper')
import rabit

rabit.init()
n = 3
rank = rabit.get_rank()
a = np.zeros(n)
for i in xrange(n):
    a[i] = rank + i

print '@node[%d] before-allreduce: a=%s' % (rank, str(a))
a = rabit.allreduce(a, rabit.MAX)
print '@node[%d] after-allreduce: a=%s' % (rank, str(a))
rabit.finalize()