title: Vectorization guidelnes for jax
body: <p>suppose I have a function (for simplicity, covariance between two  series, though the question is more general):</p>
  <pre><code>def cov(x, y):
     return jnp.dot((x-jnp.mean(x)), (y-jnp.mean(y)))
  </code></pre>
  <p>Now I have a &quot;dataframe&quot; <code>D</code> (a 2-dimenisonal array, whose columns are my series) and I want to vectorize <code>cov</code> in such a way that the application of the vectorized function produces the covariance matrix. Now, there is an obvious way of doing it:</p>
  <pre><code>cov1 = jax.vmap(cov, in_axes=(None, 1))
  cov2 = jax.vmap(cov1, in_axes=(1, None))
  </code></pre>
  <p>but that seems a little clunky. Is there a &quot;canonical&quot; way of doing this?</p>