title: Calculate only lower triangular elements of a matrix OR calculation on all possible pairs of the elements of a vector with jax
body: <p>Is it possible to efficiently run some calculation on all possible pairs of the elements of a vector? I.e. I want to fill the lower triangular elements of a matrix (possibly flattened).</p>
  <p>I.e. I want to:</p>
  <ul>
  <li>calculate <code>do_my_calculation(input_vector[i], input_vector[j])</code></li>
  <li>for all i, j in [1, <code>length(input_vector)</code>] <strong>and</strong> j &lt; i</li>
  <li>save all the results</li>
  </ul>
  <p>The shape of the results is not terribly important. If I can choose however, I would prefer a vector corresponding to an unrolled of the triangular (i, j) matrix however.</p>
  <p>To illustrate what I would like to do in pseudo-python:</p>
  <pre><code>input_vector = np.arange(100)
  result_vector = []
  
  for i in range(1, len(input_vector)):
      for j in range(0, i):
          result_vector.append(do_my_calculation(input_vector[i], input_vector[j])
  </code></pre>
  <p>Note: For this question, the types of <code>input_vector</code> and <code>result_vector</code> in the above code are not pertinent. Equally, I am of course happy to preallocate result_vector if required. I am using a list for the sake of conciseness of the sample code.</p>
  <h1>Edit 1: concrete example as requested by @ddejohn</h1>
  <p>Note: The question is not whether I can get this to run in jax but whether I can get it to run efficiently, i.e. vectorized .</p>
  <pre><code># Set up the problem
  import numpy as np
  
  dim = 15
  input_vector_x = np.random.rand(dim)
  input_vector_y = np.random.rand(dim)
  output_vector = np.empty(np.tril_indices(dim, k=-1)[0].size)
  
  assert input_vector_x.size == input_vector_y.size
  
  # alternative implementation 1
  counter = 0
  for i in range(1, input_vector_x.size):
      for j in range(0, i):
          output_vector[counter] = (input_vector_y[j] - input_vector_y[i]) / (input_vector_x[j] - input_vector_x[i])
          counter += 1
  
  # alternative implementation 2
  
  indices = np.tril_indices(dim, k=-1)
  i = indices[0]
  j = indices[1]
  
  output_vector = (input_vector_y[j] - input_vector_y[i]) / (input_vector_x[j] - input_vector_x[i])
  </code></pre>