Behavior of jnp.fft.fftn(jnp.fft.ifftn(x)) under jit compilation

[mjo22]

My understanding is the XLA compiler under jit turns jnp.exp(jnp.log(x)) into x. Does it do the same for jnp.fft.fftn(jnp.fft.ifftn(x))?

[jakevdp]

In the current version of JAX, fftn(ifftn(x)) is not fused or canceled out by the compiler; you can see this by using ahead-of-time lowering to print the compiled lower-level representation:

import jax
import jax.numpy as jnp

def g(x):
  return jnp.fft.fftn(jnp.fft.ifftn(x))
print(jax.jit(g).lower(jnp.arange(100.0)).compile().as_text())

HloModule jit_g, entry_computation_layout={(f32[100]{0})->c64[100]{0}}, allow_spmd_sharding_propagation_to_output={true}

ENTRY %main.11 (Arg_0.1: f32[100]) -> c64[100] {
  %Arg_0.1 = f32[100]{0} parameter(0), sharding={replicated}
  %convert.0 = c64[100]{0} convert(f32[100]{0} %Arg_0.1), metadata={op_name="jit(g)/jit(main)/jit(fft)/convert_element_type[new_dtype=complex64 weak_type=False]" source_file="<ipython-input-2-81ad187ab559>" source_line=2}
  %fft.0 = c64[100]{0} fft(c64[100]{0} %convert.0), fft_type=IFFT, fft_length={100}, metadata={op_name="jit(g)/jit(main)/jit(fft)/fft[fft_type=FftType.IFFT fft_lengths=(100,)]" source_file="<ipython-input-2-81ad187ab559>" source_line=2}
  ROOT %fft.1 = c64[100]{0} fft(c64[100]{0} %fft.0), fft_type=FFT, fft_length={100}, metadata={op_name="jit(g)/jit(main)/jit(fft)/fft[fft_type=FftType.FFT fft_lengths=(100,)]" source_file="<ipython-input-2-81ad187ab559>" source_line=2}
}

[mjo22]

Thank you!