Fix/isoneutral

[serge-sans-paille]

No description provided.

[serge-sans-paille]

cc @ashwinvis

[ashwinvis]

It works. Thanks!

[serge-sans-paille]

@ashwinvis and what about the speedup?

[ashwinvis]

Some numbers:

• numba implementation is a rewrite, uses loops instead of vectorized ops
• jax and pythran makes an interesting comparison

benchmarks.isoneutral_mixing
============================
Running on CPU

size          backend               calls     mean      stdev     min       25%       median    75%       max       Δ       
----------------------------------------------------------------------------------------------------------------------------
          16  numba                  100,000     0.000     0.000     0.000     0.000     0.000     0.000     0.000    43.476
          16  transonic_pythran      100,000     0.000     0.000     0.000     0.000     0.000     0.000     0.000    16.115
          16  jax                      1,000     0.000     0.000     0.000     0.000     0.000     0.000     0.000    14.237
          16  aesara                  10,000     0.000     0.000     0.000     0.000     0.000     0.000     0.001     1.847
          16  numpy                   10,000     0.001     0.000     0.001     0.001     0.001     0.001     0.001     1.000
          16  pytorch                 10,000     0.002     0.000     0.002     0.002     0.002     0.002     0.002     0.412
         128  numba                  100,000     0.000     0.000     0.000     0.000     0.000     0.000     0.000    14.784
         128  jax                      1,000     0.000     0.000     0.000     0.000     0.000     0.000     0.000     9.880
         128  transonic_pythran       10,000     0.000     0.000     0.000     0.000     0.000     0.000     0.000     6.019
         128  aesara                  10,000     0.000     0.000     0.000     0.000     0.000     0.000     0.001     1.883
         128  numpy                   10,000     0.001     0.000     0.001     0.001     0.001     0.001     0.001     1.000
         128  pytorch                 10,000     0.002     0.000     0.002     0.002     0.002     0.002     0.002     0.460
       1,024  jax                      1,000     0.000     0.000     0.000     0.000     0.000     0.000     0.000     5.432
       1,024  numba                   10,000     0.000     0.000     0.000     0.000     0.000     0.000     0.001     4.711
       1,024  transonic_pythran       10,000     0.001     0.000     0.001     0.001     0.001     0.001     0.001     2.715
       1,024  aesara                  10,000     0.001     0.000     0.001     0.001     0.001     0.001     0.002     1.857
       1,024  numpy                   10,000     0.002     0.000     0.002     0.002     0.002     0.002     0.003     1.000
       1,024  pytorch                 10,000     0.002     0.000     0.002     0.002     0.002     0.002     0.004     0.773
       8,192  jax                      1,000     0.003     0.000     0.003     0.003     0.003     0.003     0.005     2.694
       8,192  numba                    1,000     0.004     0.000     0.004     0.004     0.004     0.004     0.006     2.106
       8,192  transonic_pythran        1,000     0.005     0.000     0.005     0.005     0.005     0.005     0.009     1.539
       8,192  aesara                   1,000     0.005     0.000     0.005     0.005     0.005     0.005     0.008     1.511
       8,192  pytorch                  1,000     0.006     0.000     0.005     0.006     0.006     0.006     0.009     1.402
       8,192  numpy                    1,000     0.008     0.000     0.008     0.008     0.008     0.008     0.013     1.000