Using a global config parameters and rejits

[ASKabalan]

Was wondering if there is an easy and minimal way to reproduce this behaviour with our own config parameters

import jax
from functools import partial
from jax import numpy as jnp
@partial(jax.jit , static_argnums=(0,))
def f(shape):
    print(f"jitted for shape {shape}")
    return jnp.ones(shape)
jax.config.update("jax_enable_x64", False)
shape = (2 , 2 , 2)
f(shape)

# Out:
# jitted for shape (2, 2, 2) 
# Array([[[1., 1.],
#         [1., 1.]],
# 
#        [[1., 1.],
#         [1., 1.]]], dtype=float32)

jax.config.update("jax_enable_x64", True)

f(shape)
# Out:
# jitted for shape (2, 2, 2)
# Array([[[1., 1.],
#         [1., 1.]],
# 
#        [[1., 1.],
#         [1., 1.]]], dtype=float64)

Currently, if I change a global config parameter the function is not rejitted
As a workaround I use jax.clear_caches() but I am hoping for a cleaner solutions

Thank you

[jakevdp]

Hi - thanks for the question! The way that this works for configurations like enable_x64 is via this function:

jax/jax/_src/config.py

Lines 194 to 201 in f833891

	def trace_context():
	"""Returns a tuple of configuration values that affect tracing.
	These values are included in the cache key for linear_util.cache.
	Values included in this set should also most likely be included in
	the C++ JIT state, which is handled separately.
	"""

This is an internal utility, and unfortunately there is no public API for customizing the trace context.

[ASKabalan]

Thank you for the quick answer @jakevdp !!

Can we clear a specific entry in the jit cache? without clearning everything
I can assume that it is also possible but with no public API aswell, no?

[jakevdp]

No, as far as I know there's no mechanism (public or non-public) for clearing a single cache entry.