Very simple compilation question

[unalmis]

In the compiled function, the number of flops required for g is larger than f. I thought this was expected, but I am not sure if JAX advertises they should be the same.

from jax import jit
import jax.numpy as jnp

x = jnp.ones((1000, 1000))

f = jit(lambda x: jnp.sum(x) * 5)
print(f.lower(x).compile().cost_analysis())
g = jit(lambda x: jnp.sum(x * 5))
print(g.lower(x).compile().cost_analysis())

[jakevdp]

I mentioned this in #31257 as well, but I suspect this is a deliberately skipped optimization, because there are cases where changing the order of operations in this way would lead to overflow. Here's a simple example:

In [1]: import jax.numpy as jnp

In [2]: a = jnp.float32(1E-4)

In [3]: x = 1E37 * jnp.arange(10)

In [4]: (a * x).sum()
Out[4]: Array(4.5e+34, dtype=float32)

In [5]: a * x.sum()
Out[5]: Array(inf, dtype=float32)

If the compiler were to automatically rewrite (a * x).sum() to a * x.sum(), then JIT-compiling a program with these values would cause it to overflow.

Note that the source of truth for these kinds of compiler decisions is not in JAX, but rather in https://github.com/openxla/xla – so folks at that repository might know more about this type of optimization, and whether there might be compiler flags that would enable it.

[unalmis]

For any function, one can always engineer an equivalent formula to compute it such that it will lead to overflow in finite precision unless the instructions are evaluated in a particular order. e.g. $f(x) = 2 x$ can be implemented as return x + 1e58 - 1e57 + 1e52 - 1e52 + 1e57 - 1e58 + x. Examples similar in spirit to this will overflow unless the compiler recognizes the intermediate operations as an identity.

Therefore, I do not understand the purpose of avoiding the compiler optimization in the original issue because it is an attempt to attain the impossible goal of avoiding finite precision error.

One can just as easily generate an example where not performing the optimization to pull the scalar out of the sum leads to overflow. E.g if jnp.sum(x) = 0 in finite precision but x was an array of large oscillating magnitude.

[unalmis]

In general, finite precision error is reduced by reducing the number of floating point operations. Hence the better default is to always pull the scalar out of the sum to reduce flops.

[jakevdp]

I think you're misunderstanding me: the goal of the compiler is not to avoid overflow, or improve floating point error accumulation; the goal of the compiler is to optimize code without significantly affecting the numerics as expressed in the original program. If the original program does not overflow, the compiled program should maintain that property. If the original program does overflow, the compiled program should maintain that property.

Of course, floating point math being what it is, you can never guarantee exact bitwise equivalence before and after a compiler rewrite, but what you can do is avoid certain optimizations that have been found to be problematic in practice (and XLA has removed certain optimizations over the years for this reason).

You may disagree with the choices that the compiler engineers have made, and that's fine – this is an area in which reasonable people can disagree. I'll note that often there are flags to enable "unsafe" optimization passes for users who want different behavior (the --xla_enable_fast_math flag is an example), so you might explore whether any of those is relevant to your own work. The choice of which optimizations to enable by default is not made in this repository, though, so if you want to discuss things at that level I'd suggest engaging with folks at https://github.com/openxla/xla/.

[unalmis]

Optimizing a reduction operation toward using less flops and reducing floating point error in general is no different than optimizing log(exp) as an identity. log(exp(1e30)) overflows in non-compiled code but the compiler changes behavior to avoid overflow. So I am not convinced that "preserving overflow" is a policy that XLA attempts to maintain. Recognizing identity operations and avoiding them has always been a goal however.

Yes the original discussion should have been posted to XLA, but given that there was a response here I replied. Thanks for the input.

[jakevdp]

Recognizing identity operations and avoiding them has always been a goal however.

Yes, but that goal is constrained by the additional goal of avoiding such reductions when it causes problems downstream. e.g. lambda x, y: x + y - x previously was reduced to lambda x, y: y, but it no longer is because that optimization caused numerical problems in important models. Similarly, exp(log(x)) I believe is not reduced to an identity, but rather to a log-exp fusion.

I don't think it's helpful to continue arguing here about which optimizations should be enabled by default, because that's a decision made by XLA, not by JAX.