_cache_size() changes in 0.3.15 compared to 0.3.10?

[pipme]

import jax
import jax.numpy as jnp

@jax.jit
def f(Y):
    return jnp.sum(Y**2)

Y = jnp.ones(2)

for i in range(5):
    print(f._cache_size())
    f(Y)

The above prints 0 1 2 2 2 in Jax 0.3.15 while 0 1 1 1 1 in Jax 0.3.10. It seems related to ** 2. Why that happens? Thanks.

[jakevdp]

Thanks for the question - I'm not sure what might have caused this change. Doing some investigation, it seems to not be related to the jaxlib version, but rather to the jax version, and changed sometime between jax 0.3.14 and 0.3.15. I modified your script to print the versions:

import jax
import jaxlib
import jax.numpy as jnp

print(f"jax: {jax.__version__}")
print(f"jaxlib: {jaxlib.__version__}")

@jax.jit
def f(Y):
    return jnp.sum(Y**2)

Y = jnp.ones(2)

for i in range(5):
    print(f._cache_size(), end=' ')
    f(Y)
print()

Here are the outputs of two runs:

jax: 0.3.14
jaxlib: 0.3.14
0 1 1 1 1 

jax: 0.3.15
jaxlib: 0.3.14
0 1 2 2 2 

Among changes to jax/_src/api.py in the date range between jax 0.3.14 and 0.3.15, the only one that looks like it could have affected this is #11399, though that's just a guess. In any case, I doubt this change was intentional.

[jakevdp]

I suspected wrong - after bisecting, I found that this PR caused the change in behavior: #11298

With jaxlib 0.3.14, prior to this PR we see 0 1 1 1 1, and after this PR we see 0 1 2 2 2, due probably to the axes environment changing between the first and second call. I'm not certain whether this is expected; perhaps @pschuh could weigh in?

[jakevdp]

Ah, found the culprit. In this line: https://github.com/google/jax/blob/70bc0da99517dec729f84eaf6edf27954fe9e462/jax/_src/config.py#L491

The state defaults to None. If we instead default to (), then we avoid the extra cache entry.

[jakevdp]

#11883 has the fix. Thanks for the report!

[pipme]

Cool thanks!