title: Python: how to write this code to run on GPU?
body: <p>I have been trying for quite some time to implement my code to run on GPU, however with little success. I would really appreciate someone helping with the implementation.</p>
  
  <p>Let me say a few words about the problem. I have a graph G with N nodes and a distribution mx on each node x. I would like to compute the distance between the distributions for every pair of nodes for all edges. For a given pair, (x,y), I use the code <code>ot.sinkhorn(mx, my, dNxNy)</code> from the python POT package to compute the distance. Again, mx, my are vectors of size Nx and Ny on nodes x and y and dNxNy is a Nx x Ny distance matrix.</p>
  
  <p>Now, I discovered that there is a GPU implementation of this code <code>ot.gpu.sinkhorn(mx, my, dNxNy)</code>. However, this is not good enough because I mx, my and dNxNy would need to be uploaded to the GPU at every iteration, which is a massive overhead. So, the idea is to parallelise this for all edges on GPU.</p>
  
  <p>The essence of the code is as follows. mx_all is all the distributions </p>
  
  <pre><code>for i,e in enumerate(G.edges):
      W[i] = W_comp(mx_all,dist,e)
  
  def W_comp(mx_all, dist,  e):
      i = e[0]
      j = e[1]
  
      Nx = np.array(mx_all[i][1]).flatten()
      Ny = np.array(mx_all[j][1]).flatten()
      mx = np.array(mx_all[i][0]).flatten()
      my = np.array(mx_all[j][0]).flatten()
  
      dNxNy = dist[Nx,:][:,Ny].copy(order='C')
  
      W = ot.sinkhorn2(mx, my, dNxNy, 1)
  </code></pre>
  
  <p>Below is a minimal working example. Please ignore everything except the part between dashed <code>===</code> signs. </p>
  
  <pre><code>import ot
  import numpy as np
  import scipy as sc
  
  
  def main():
      import networkx as nx
  
      #some example graph
      G = nx.planted_partition_graph(4, 20, 0.6, 0.3, seed=2)
      L = nx.normalized_laplacian_matrix(G)
  
      #this just computes all distributions (IGNORE)
      mx_all = []
      for i in G.nodes:
          mx_all.append(mx_comp(L,1,1,i))  
  
      #some random distance matrix (IGNORE)
      dist = np.random.randint(5,size=(nx.number_of_nodes(G),nx.number_of_nodes(G)))          
  
  # ============================================================================= 
  #this is what needs to be parallelised on GPU
      W = np.zeros(nx.Graph.size(G))
      for i,e in enumerate(G.edges):
          print(i)
          W[i] = W_comp(mx_all,dist,e)
  
      return W
  
  def W_comp(mx_all, dist,  e):
      i = e[0]
      j = e[1]
  
      Nx = np.array(mx_all[i][1]).flatten()
      Ny = np.array(mx_all[j][1]).flatten()
      mx = np.array(mx_all[i][0]).flatten()
      my = np.array(mx_all[j][0]).flatten()
  
      dNxNy = dist[Nx,:][:,Ny].copy(order='C')
  
      return ot.sinkhorn2(mx, my, dNxNy,1)
  
  # =============================================================================
  
  #some other functions (IGNORE)
  def delta(i, n):
  
      p0 = np.zeros(n)
      p0[i] = 1.
  
      return p0
  
  # all neighbourhood densities
  def mx_comp(L, t, cutoff, i):
      N = np.shape(L)[0]
  
      mx_all = sc.sparse.linalg.expm_multiply(-t*L, delta(i, N))
      Nx_all = np.argwhere(mx_all &gt; (1-cutoff)*np.max(mx_all))
  
      return mx_all, Nx_all  
  
  if __name__ == "__main__":
      main()  
  </code></pre>
  
  <p>Thank you!!</p>