chex.variants(with_pmap=True) ignores `static_argnames`

The [`_with_pmap`](https://github.com/deepmind/chex/blob/70350fd8fb0937034c8da7fd1dd47de7aad0747a/chex/_src/variants.py#L428) function accepts `static_argnums` as a parameter, but not `static_argnames`. This is inconsistent with other variants, such as [`with_jit`](https://github.com/deepmind/chex/blob/70350fd8fb0937034c8da7fd1dd47de7aad0747a/chex/_src/variants.py#L346) and [`with_device`](https://github.com/deepmind/chex/blob/70350fd8fb0937034c8da7fd1dd47de7aad0747a/chex/_src/variants.py#L376). Crucially, this prevents to test methods that require to pass arguments by name (e.g., Distrax's [Distribution.sample()](https://github.com/deepmind/distrax/blob/af9d16336bc02e137d1892efc875c84ecb3b9361/distrax/_src/distributions/distribution.py#L115-L118))

More generally, it would be best if all variants accepted the same parameters where possible (i.e., where not specific to a single variant) and I would suggest to check all keys in `**unused_kwargs` against a list of allowed parameters (i.e., the union of the parameters of all variant functions) to prevent silent errors due to e.g., misspells.