`sharding` parameter bug in JAX 0.4.31

[balancap]

Some additional sharding parameter is passed to primitives in JAX 0.4.31, creating issues in Scalify tracer. We need a quick patch, as well as a proper long term strategy to deal with sharding in JAX.