JIT'ed JNP FFT speed on TPU and lax scan on function with 3 input args

[sarahbton]

Hi Jax Team,

Thanks for the great package. I have a couple questions I have run into while implementing an algorithm that I was hoping a JAX expert could help me out with.

import jax.tools.colab_tpu
jax.tools.colab_tpu.setup_tpu()
from jax import random
import jax 
import jax.numpy as jnp
import numpy as np 
key = random.PRNGKey(0)
key, subkey = random.split(key)
x = random.normal(key, (5000, 5000))
x_cpu = np.array(x)
%timeit -n 5 -r 5 np.fft.fft(x_cpu) 5 loops, best of 5: 587 ms per loop
%timeit -n 5 -r 5 jnp.fft.fft(x).block_until_ready() #5 loops, best of 5: 250 ms per loop

When using PMAP we get down to around 250ms per loop which is obviously around 8x faster than non-PMAP'ed.

from jax import pmap
keys = random.split(random.PRNGKey(0), 8)
mats = pmap(lambda key: random.normal(key, (5000, 5000)))(keys)
%timeit -n 5 -r 5 pmap(jnp.fft.fft)(mats).block_until_ready() #5 loops, best of 5: 251 ms per loop

Yet after JIT'ing the fft the speed doesn't change at all:

from jax import jit
def fft(valu):
    return jnp.fft.fft(valu)
c = jit(fft)
%timeit -n 5 -r 5 c(x).block_until_ready() #5 loops, best of 5: 249 ms per loop
%timeit -n 5 -r 5 fft(x).block_until_ready() #5 loops, best of 5: 250 ms per loop

Unless I am mis understanding JIT, wouldn't the speed increase by more than 1ms when using the JIT'ed version of the function?

Are there any other JAX elements (not JIT or PMAP) that I should look at when trying to implement huge FFT's across TPUs?

Finally,

I am having some trouble with the lax scan function when my scanned function needs access to multiple arrays.

For example:

a = [5, 5, 5]
b = [5, 5, 5]

def myfunc(carry, x):
    global a, b
    temp = a[x] * b[x]
    carry = np.floor(a[x]/temp)
    return carry, temp
carry = 0
for i in range(0, 3):
    carry, x = myfunc(carry, i)
    print(carry, x)

I tried converting the above into a lax scan by passing the a and b as tuples along with the i value, but it continued to throw errors. Is there a correct way to implement something like that (that's hopefully computationally efficient)?

-Sarah

[jakevdp]

For your first question, why wrapping a single fft operation with jit does not make it faster: this is because fft is jit-compiled by default: https://github.com/google/jax/blob/2f7e125c70558275554015209a64d59baf913b70/jax/_src/lax/fft.py#L46-L47

If you're using the FFT within a sequence of other operations, JIT-compiling the full sequence of operations may yield improvement, but re-compiling an already compiled operation won't yield any speedups.

For your second question: the problem is you're trying to index into non-jax arrays using traced indices. Try using this:

a = jnp.array([5, 5, 5])
b = jnp.array([5, 5, 5])

then you should be able to re-express this in terms of fori_loop or scan.