You might be interested in Legate [1]. It supports the NumPy interface as a drop-in replacement, supports GPUs and also distributed machines. And you can see for yourself their performance results; they're not far off from hand-tuned MPI.

[1]: https://github.com/nv-legate/legate.numpy

Disclaimer: I work on the library Legate uses for distributed computing, but otherwise have no connection.

Interesting find about the indexing. I just had the opposite experience, swapped from numpy to torch in a project and got 2000x speedup on some indexing and basic maths wrapped in autodiff. And I haven't moved it onto cuda yet.

Here's an example that illustrates the phenomenon. If memory serves me right, index latency is superlinear in dimension count.

   import time, torch
   from itertools import product

   N = 100

   ten = torch.randn(N,N,N)
   arr = ten.numpy()

   def indexTimer(val):
       start = time.time()
       for i,j,k in product(range(N), range(N), range(N)):
           x = val[i, j, k]
       end = time.time()
       print('{:.2f}'.format(end-start))

   indexTimer(ten)
   indexTimer(arr)

Ah, I'm guessing it's the loop that kills you. Arrays/tensors are supposed to be used as a whole, I imagine e.g. x=val**2 would be much faster.

I'm indexing one array with another e.g. x=y[z] to pull out all the values I want at once into another array for processing. And not using a python loop for the expensive part.

>>> built an options pricing engine on top of PyTorch

I'd love to hear more about this! Do you have any posts or write-ups on this?