title: JAX(XLA) vs Numba(LLVM) Reduction
body: <p>Is it possible to make CPU only reductions with JAX comparable to Numba in terms of computation time?</p>
  <p>The compilers come straight from <code>conda</code>:</p>
  <pre><code>$ conda install -c conda-forge numba jax
  </code></pre>
  <p>Here is a 1-d NumPy array example</p>
  <pre><code>import numpy as np
  import numba as nb
  import jax as jx
  
  @nb.njit
  def reduce_1d_njit_serial(x):
      s = 0
      for xi in x:
          s += xi
      return s
  
  @jx.jit
  def reduce_1d_jax_serial(x):
      s = 0
      for xi in x:
          s += xi
      return s
  
  N = 2**10
  a = np.random.randn(N)
  </code></pre>
  <p>Using <code>timeit</code> on the following</p>
  <ol>
  <li><code>np.add.reduce(a)</code> gives <code>1.99 µs ...</code></li>
  <li><code>reduce_1d_njit_serial(a)</code> gives <code>1.43 µs ...</code></li>
  <li><code>reduce_1d_jax_serial(a).item()</code> gives <code>23.5 µs ...</code></li>
  </ol>
  <p>Note that <code>jx.numpy.sum(a)</code> and using <code>jx.lax.fori_loop</code> gives comparable (marginally slower) comp. times to <code>reduce_1d_jax_serial</code>.</p>
  <p>It seems there is a better way to craft the reduction for XLA.</p>
  <p><strong>EDIT</strong>: compile times were not included as a print statement proceeded to check results.</p>