shard_map should support static_argnums about jax

[reinerp]

I think that's probably sufficient, thank you!

[jakevdp]

Assigning @mattjj as this is similar to #17461

[mattjj]

Hey @reinerp !

What if instead of static_argnums, we allowed passing in_specs=(P(...), P(...), None, None)? That'd be more like vmap's in_axes=None. (Or do you prefer static_argnums? In that case, what do you pass for the corresponding in_specs? Or do we just want to leave out the corresponding entries of in_specs?)

[mattjj]

Assuming using in_specs=(..., None, ...) rather than static_argnums=... is okay, #22049 should fix this!