Does the compiler ignore computations that are multiplied by zero?

[Rasmuskh]

Hi,

I have two questions:

(1) Does the Jax compiler know to not carry out calculations which are multiplied by zero?
If for example I have a function $f(k, A, B) = k AB $, then if the scalar $k!=0$ it is necessary to compute the whole right hand side, but if $k=0$, then it is sufficient to just check the dimensions of the matrices A and B and return an array of zeros with correct dimensions. Does the compiler perform this kind of optimization?

(2) What kind of introspection tools are there? Is there a way I can see the low level code generated by the compiler?

Thanks in advance! :)

[jakevdp]

In general, no it does not, because multiplication by zero cannot safely be ignored. Consider this function:

def f(a, b):
  return a + 0 * b

def f_ignore_mult_by_zero(a, b):
  return a

When b is finite, the functions return the same results:

f(1.0, 1.0) # 1.0
f_ignore_mult_by_zero(1.0, 1.0) # 1.0

However, when b is NaN or inf, the functions return different results:

f(1.0, np.inf) # NaN
f_ignore_mult_by_zero(1.0, np.inf) # 1.0

One of the key goals of the compiler is to not optimize in a way that will change program outputs, and so the compiler will not elide multiplication by zero.

If you want "ignore multiplication by zero" semantics, you can do this explicitly with lax.cond; e.g. something like this:

def f(a, b, k):
  return a + lax.cond(k == 0, lambda: 0.0, lambda: k * b)

[Rasmuskh]

Thanks,
Good point regarding NaN and Inf.
I will try to speed my code up with lax.cond :)

Regarding my second question, then I am wondering if something like julias @code_native and @code_llvm exists for jax?

[jakevdp]

Sorry, I missed the second question. To see the optimized HLO emitted by the compiler, you can use the tools described at https://jax.readthedocs.io/en/latest/aot.html
For a function with a call signature func(*args) (where args are arrays or other traceable variables) it would look like this:

print(jax.jit(func).lower(*args).compile().as_text())

[Rasmuskh]

Thanks!