don't use tree_map for pjit arg checking
don't use tree_map for pjit arg checking

We only need to get the argument leaves for error checking. By avoiding
tree_map we avoid creating (and then throwing away) a results tree with
Nones at the leaves (since `_check_arg` returns None). More importantly,
this makes the code work with pytrees which don't actually support
containing Nones (like the PRNGKeyArray pytree!).

google/jax,google/jax,google/jax,google/jax