Tell jit to assume that x * 0 = 0

[mariogeiger]

Dear jax developers,

If I print the xla output of x + 0.0 I get

tmp_0 = parameter(0)
tmp_1 = copy(tmp_0)
ROOT tmp_2 = tuple(tmp_1)

If I print the xla output of x * 0.0 I get

tmp_0 = parameter(0)
tmp_1 = fusion(tmp_0), kind=kLoop, calls=
(param_0: f32[3]) -> f32[3] {
  tmp_0 = parameter(0)
  tmp_1 = constant(0)
  tmp_2 = broadcast(tmp_1), dimensions={}
  ROOT tmp_3 = multiply(tmp_0, tmp_2)
}
ROOT tmp_2 = tuple(tmp_1)

My understanding is that it is due to the fact that nan * 0.0 and inf * 0.0 are not equal to zero.

Is there a way to gently ask the compiler to assume that x * 0.0 = 0.0 for specific parts of the code?

Is there an unsafe_multiplication operation and a function transformation that replaces all multiplication by their unsafe counterparts?

Best regards,
Mario

[jakevdp]

I suspect this is not necessarily due to the need to handle NaN and Inf values, but because 0 is not an identity for multiplication, and so there's no special rule for it. If you generate code for x * 1 rather than x * 0, I suspect you'll see the same shortcut used in the x + 0 case.

[mariogeiger]

I agree with you that 0 is not the identity for multiplication but I'm still quite convinced that the axioms of a field imply that x * 0 = 0.

Yes, x * 1 = 1 is properly simplified by jax.jit.

[jakevdp]

Sure, but floating point numbers do not form a field. I am not sure there is any method to tell XLA to do the kind of unsafe multiplication that you have in mind.

[nouiz]

Note, this optimization is done for integer. In case you can change the dtype.

[hawkinsp]

BTW, there is a mechanism to tell JAX to do this optimization: it's "fast math".

Try setting flags like XLA_FLAGS="--xla_cpu_fast_math_honor_infs=false --xla_cpu_fast_math_honor_nans=false". I'm not sure if that will work, but some combination of the XLA fast math flags should do the trick.