shard_map should support static_argnums

[reinerp]

We've started using (a type-based derivative of) shard_map as a function decorator. Ideally we'd like to annotate our top-level training function with @jax.jit @typed_shard_map decorator on a single top-level function. However, shard_map doesn't support static_argnums. We're forced to add a level of nested functions, in order to capture the static arguments in a closure, see this example:

https://github.com/MatX-inc/seqax/blob/871bd91efc004e569aa91228ef0a66931ed16986/train.py#L292-L295

Would it be possible to add static_argnums support to shard_map? That would allow us to simplify the above to:

@partial(jax.jit, static_argnums=(2, 3), donate_argnums=(0,))
@typed_shard_map
def training_step(state: State, step: u32[b''], h: Hparams, hparams: TrainingHparams, batch: TokenBatch) -> Tuple[Any, Metrics]:
  ...

[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!

[reinerp]

I think that's probably sufficient, thank you!