jit lower depends on process trace history #40312

import re
import jax, jax.numpy as jnp
from jax import lax
from jax.experimental import topologies

jax.config.update("jax_default_device", topologies.get_topology_desc("v5e:2x2").devices[0])
x_spec = jax.ShapeDtypeStruct((16, 128), jnp.float32)
y_spec = jax.ShapeDtypeStruct((16,), jnp.int32)

@jax.jit
def foo(y):
    def body(i, carry):
        foos, acc = carry
        return foos.at[i].set(acc), acc + y[i]
    return lax.fori_loop(0, y.shape[0], body, (jnp.zeros_like(y), jnp.int32(0)))

@jax.jit
def bar(x, y):
    foos, acc = foo(y)
    return x * foos.astype(jnp.float32)[:, None], acc.astype(jnp.float32)

def f(x, y):
    for _ in range(2):
        baz, acc = bar(x, y)
        x = baz + acc
    baz, _ = bar(x, y)
    return baz * 5.0

def g(x, y):
    baz, acc = bar(x, y)
    return baz - acc

def lower(fn):
    return jax.jit(lambda x, y: fn(x, y)).lower(x_spec, y_spec)

keep = []
def lower_unrelated(n):
    for _ in range(n):
        inner = jax.jit(lambda z: z + 1.0)
        outer = jax.jit(lambda z: inner(z) * 2.0)
        keep.append(outer); outer.lower(x_spec)

def report(label, lowered):
    n_funcs = len(re.findall(r"func\.func ", lowered.as_text()))
    n_inst = sum(1 for l in lowered.compile().as_text().splitlines() if l.startswith("  ") and " = " in l)
    print(f"{label:22s} StableHLO functions={n_funcs}  optimized HLO instructions={n_inst}")

lower(g); lower_unrelated(1000); lower(g); lower_unrelated(1500)

report("after history:", lower(f))         
# StableHLO functions=7  optimized HLO instructions=67

jax.clear_caches()

report("after clear_caches():", lower(f))  
# StableHLO functions=6  optimized HLO instructions=76