Does JAX's XLA compiler optimize away sqrt followed by square?

[NeilGirdhar]

x = jnp.array(1+1j)
y  = jnp.square(jnp.abs(x))

[hawkinsp]

Here's a semi-internal way to look at the compiled HLO:

import jax, jax.numpy as jnp

def f(x):
 return jnp.square(jnp.abs(x))

x = jnp.array(1+1j)
c = jax.xla_computation(f)(x)

backend = jax.lib.xla_bridge.get_backend()
e = backend.compile(c)
print(e.hlo_modules()[0].to_string())

Output on CPU:

WARNING:absl:No GPU/TPU found, falling back to CPU. (Set TF_CPP_MIN_LOG_LEVEL=0 and rerun for more info.)
HloModule xla_computation_f.6

%fused_computation (param_0.1: c64[]) -> f32[] {
  %param_0.1 = c64[] parameter(0)
  %abs.0 = f32[] abs(c64[] %param_0.1), metadata={op_type="abs" op_name="xla_computation(f)/abs" source_file="/Users/phawkins/p/jax/t.py" source_line=4}
  ROOT %multiply.0 = f32[] multiply(f32[] %abs.0, f32[] %abs.0), metadata={op_type="integer_pow" op_name="xla_computation(f)/integer_pow[ y=2 ]" source_file="/Users/phawkins/p/jax/t.py" source_line=4}
}

ENTRY %xla_computation_f.6 (parameter.1: c64[]) -> (f32[]) {
  %parameter.1 = c64[] parameter(0)
  %fusion = f32[] fusion(c64[] %parameter.1), kind=kLoop, calls=%fused_computation, metadata={op_type="integer_pow" op_name="xla_computation(f)/integer_pow[ y=2 ]" source_file="/Users/phawkins/p/jax/t.py" source_line=4}
  ROOT %tuple.5 = (f32[]) tuple(f32[] %fusion)
}

So no, it looks like XLA left both operations alone, at least at this level of the compiler.

[NeilGirdhar]

Okay, thanks for showing me this trick (I think it was recently shown to me in another issue too).

I guess I should keep using my abs_square function. 😄