XLA change significantly increase the computational time

Description

The XLA change in openxla/xla@d97f3a7bffe359846e45948e2cab8a6c9be2993a (pulled into JAX with commit 54ba3e02012afe6663943455ce57fea0f8d2cddd) increased the compilation time for the example below (using the JAX-based variPEPS/variPEPS_Python@465915c22ed1a17f14cb17c6e28ca4985ee13bd7 library) by a factor of ~5.

At the time of commit 54ba3e02012afe6663943455ce57fea0f8d2cddd it is still possible to mitigate the increase of the compilation time with XLA_FLAGS='--xla_cpu_experimental_enable_tiling_propagation=false'. This mitigation is not possible with the XLA change in openxla/xla@8d14527c7fddf08f994d48bd696a84db26064394 (pulled into JAX with commit 4c7da2781a48620f33386f1084ab9069b80f03b9). After that change the xla flag does not effect the compilation time anymore.

Example:

import time

import varipeps
import jax.numpy as jnp

Id = jnp.eye(2)
Sx = jnp.array([[0, 1], [1, 0]]) / 2
Sy = jnp.array([[0, -1j], [1j, 0]]) / 2
Sz = jnp.array([[1, 0], [0, -1]]) / 2

gate = (jnp.kron(Sx, Sx) + jnp.kron(Sy, Sy) + jnp.kron(Sz, Sz))

varipeps.config.config.ctmrg_print_steps = True
varipeps.config.config.ctmrg_full_projector_method = varipeps.config.Projector_Method.FULL_QR
varipeps.config.ctmrg_convergence_eps = 1 * 1e-7
varipeps.config.ad_custom_print_steps = True
varipeps.config.ad_custom_fixed_point_method = varipeps.config.Grad_Fixed_Point_Method.ITERATIVE
varipeps.config.ctmrg_heuristic_decrease_chi = False
varipeps.config.ctmrg_increase_truncation_eps = False
varipeps.config.ctmrg_heuristic_increase_chi = False

unitcell = varipeps.peps.PEPS_Unit_Cell.random(
    ((0, 1), (1, 0)),
    2,
    2,
    20,
    float,
    20,
    varipeps.peps.PEPS_Type.SQUARE,
    seed=6396316583739234,
)

# Measure execution time including jitting
start = time.perf_counter()
result, _ = varipeps.ctmrg.calc_ctmrg_env(tuple(i.tensor for i in unitcell.get_unique_tensors()), unitcell, enforce_elementwise_convergence=False)
print(time.perf_counter() - start)

System info (python version, jaxlib version, accelerator, etc.)

jax:    0.11.1.dev20260726+4c7da2781a
jaxlib: 0.11.1.dev0+selfbuilt
numpy:  2.5.2
python: 3.14.5 (main, May 10 2026, 19:28:16) [Clang 22.1.3 ]
device info: cpu-1, 1 local devices"
process_count: 1
platform: uname_result(system='Linux', node='qmio20', release='6.12.101+deb13-amd64', version='#1 SMP PREEMPT_DYNAMIC Debian 6.12.101-1 (2026-08-05)', machine='x86_64')