Add static_argnames support in jax.checkpoint, as in jax.jit

[oliverdutton-iso]

Hi,

jax.checkpoint is brilliant. We normally use static_argnums for things we force to be key word arguments so we are currently having to write wrappers for modules accepting static things passed as kwargs. Please could support for static_argnames be added, as in jax.jit?

e.g. in FLAX Linen (which is just lifting and passing to jax.checkpoint):

def checkpoint(
    mdl,
    *args,
    static_argnames=('is_training',),
    **kwargs,
):
  """Checkpoint a module."""
  static_kwargs = {k: v for k, v in kwargs.items() if k in static_argnames}
  other_kwargs = {k: v for k, v in kwargs.items() if k not in static_argnames}

  @functools.partial(nn.checkpoint, static_argnums=(1,))
  def f(mdl, static_args, *args, **kwargs):
    return mdl(
        *args, **kwargs, **dict(zip(static_kwargs.keys(), static_args))
    )
  return f(mdl, tuple(static_kwargs.values()), *args, **other_kwargs)

[jakevdp]

Related to this proposal: #10614