Whether the hlo generated by jax through jax.fit.lower() is optimized by xla

[songh11]

I see from the documentation that hlo is input to the xla compiler, so is hlo not XLA-optimized?
I don't know much about this and I hope someone can help with that.

[jakevdp]

jit(f).lower(...) returns un-optimized HLO. If you call jit(f).lower(...).compile(), it will be passed through the compiler for optimization. Here's a simple example:

import jax

def f(x):
  return jax.lax.sin(x) + jax.lax.sin(x)

lowered = jax.jit(f).lower(1.0)

print(lowered.as_text())

module @jit_f {
  func.func public @main(%arg0: tensor<f32> {jax.arg_info = "x", mhlo.sharding = "{replicated}"}) -> (tensor<f32> {jax.result_info = ""}) {
    %0 = stablehlo.sine %arg0 : tensor<f32>
    %1 = stablehlo.sine %arg0 : tensor<f32>
    %2 = stablehlo.add %0, %1 : tensor<f32>
    return %2 : tensor<f32>
  }
}

Notice that there are two duplicate sine computations. Now if you compile it you get this:

print(lowered.compile().as_text())

HloModule jit_f, entry_computation_layout={(f32[])->f32[]}, allow_spmd_sharding_propagation_to_output={true}

%fused_computation (param_0.1: f32[]) -> f32[] {
  %param_0.1 = f32[] parameter(0)
  %sine.0 = f32[] sine(f32[] %param_0.1), metadata={op_name="jit(f)/jit(main)/sin" source_file="<ipython-input-4-1e0b387820fd>" source_line=4}
  ROOT %add.0 = f32[] add(f32[] %sine.0, f32[] %sine.0), metadata={op_name="jit(f)/jit(main)/add" source_file="<ipython-input-4-1e0b387820fd>" source_line=4}
}

ENTRY %main.5 (Arg_0.1: f32[]) -> f32[] {
  %Arg_0.1 = f32[] parameter(0), sharding={replicated}
  ROOT %fusion = f32[] fusion(f32[] %Arg_0.1), kind=kLoop, calls=%fused_computation, metadata={op_name="jit(f)/jit(main)/add" source_file="<ipython-input-4-1e0b387820fd>" source_line=4}
}

Notice that the compiler has de-duplicated the two sine computations.