jax

097222b8 - [jax] Simplify test_closed_over_constant_sharded_and_cache_hit in pjit_test

[jax] Simplify test_closed_over_constant_sharded_and_cache_hit in pjit_test Remove mock spy on ExecuteReplicated.__call__ from test_closed_over_constant_sharded_and_cache_hit. This mocking is too complex for the little test coverage it adds. Testing pjit with closed-over constants directly validates cache misses (cpp=1), output values, sharding, and fastpath cache hits (cpp=0) cleanly without monkey-patching internal runtime methods. The only thing we are dropping from coverage is that the arguments are passed properly sharded. PiperOrigin-RevId: 1004464715