Add numpy array like operations to DistributedMatrix

[zuoshifan]

Do you have some interests to provide the DistributedMatrix class with some numpy array like operations, such as slicing, so one could easily do the following calculations but leaves all the tricky details of MPI and distribution behand?

a = np.arange(10, dtype=np.float64)
A = core.DistributedMatrix([10, 20], dtype=a.dtype, block_shape=[2, 3])
A[0, :10] = a

b = np.zeros(10, dtype=np.float64)
b = A[0, :10]
assert np.allclose(a, b)

[jrs65]

Good question, I've wondered about allowing this kind of stuff a little more, though I hadn't thought in particular about the slicing aspect. There's a complicated design question of how much do we want the DistributedMatrix class to look like a numpy array, so do we want to support these kind of operations.

In the case of the slicing as you've written it, I guess it could be implemented pretty easily with calls to pxgemr2d, similar to DistributedMatrix.redistribute.

Other things I've wondered about would be:

• Supporting the various standard elementwise operators: +, -, *, /, **. This could be done pretty easily by just overloading the operators.
• Adding support for numpy ufuncs (e.g. np.exp). This may require making DistributedMatrix a numpy subclass which would be much more involved.

Does anyone have any opinions on this idea in general, or these three specific parts? @zuoshifan @kiyo-masui @kmsmith137