flax

963f9ba9 - Fix NNX jit static args with in_shardings issue #4989

Fix NNX jit static args with in_shardings issue #4989 - Add proper static argument resolution using api_util.resolve_argnums - Consolidate 10 individual test cases into 2 parametrized tests - Add StateSharding test for static arguments - Ensure in_shardings correctly handles static argument positions