Slow XLA compile time

[giordano]

We're having some slow compilation times:

2025-03-19 07:30:02.045318: E external/xla/xla/service/slow_operation_alarm.cc:73] 
********************************
[Compiling module reactant_loop_] Very slow compile? If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.
********************************
2025-03-19 07:30:14.895694: E external/xla/xla/service/slow_operation_alarm.cc:140] The operation took 2m12.850942799s

********************************
[Compiling module reactant_loop_] Very slow compile? If you want to file a bug, run with envvar XLA_FLAGS=--xla_dump_to=/tmp/foo and attach the results.
********************************

Here is the XLA dump for CPU code and for the corresponding GPU code.