Debugging unnecessary recompilation caused by GSPMDSharding({replicated})

[milutter]

I have the problem that a function re-compiles for the first 2-3 times, when I call the function in a for loop for n >> 10 times. Using the jax debug logs, I can narrow it down to differences in 'argument mapping'

jax._src.interpreters.pxla - DEBUG - Compiling [function name] for with global shapes and types [ShapedArray(uint32[2]), 
ShapedArray(float32[5]), ShapedArray(float32[5])]. Argument mapping: [UnspecifiedValue, UnspecifiedValue, UnspecifiedValue].


jax._src.interpreters.pxla - DEBUG - Compiling [function name] for with global shapes and types [ShapedArray(uint32[2]), 
ShapedArray(float32[5]), ShapedArray(float32[5])]. Argument mapping: [UnspecifiedValue, GSPMDSharding({replicated}), 
GSPMDSharding({replicated})].

Is there more documentation what the difference between UnspecifiedValue and GSPMDSharding({replicated}) is?
How can one alter or prevent this recompilation from happening?
These parameters that change mapping belong to a flax network state.