improve checkpoint / remat concreteness error with static_argnums

[mattjj]

Follow-up to #24511.

from functools import partial
import jax

@partial(jax.remat, static_argnums=(1,))
def f(x, _):
  if x > 0:
    return x
  else:
    return jnp.sin(x)

f(3., 1.)

Before:

Traceback (most recent call last):
  File "/usr/local/google/home/mattjj/packages/jax/froystig.py", line 11, in <module>
    f(3., 1.)
  File "/usr/local/google/home/mattjj/packages/jax/froystig.py", line 6, in f
    if x > 0:
       ^^^^^
jax.errors.TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[].
The error occurred while tracing the function new_fun at /usr/local/google/home/mattjj/packages/jax/jax/_src/ad_checkpoint.py:394 for checkpoint. This concrete value was not available in Python because it depends on the value of the argument dyn_args[0].
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerBoolConversionError

Consider using the `static_argnums` parameter for `jax.remat` or `jax.checkpoint`. See the `jax.checkpoint` docstring and its example involving `static_argnums`:
https://jax.readthedocs.io/en/latest/_autosummary/jax.checkpoint.html

Notice how it says "This concrete value was not available in Python because it depends on the value of the argument dyn_args[0].". What!?

After:

Traceback (most recent call last):
  File "/usr/local/google/home/mattjj/packages/jax/froystig.py", line 11, in <module>
    f(3., 1.)
  File "/usr/local/google/home/mattjj/packages/jax/froystig.py", line 6, in f
    if x > 0:
       ^^^^^
jax.errors.TracerBoolConversionError: Attempted boolean conversion of traced array with shape bool[].
The error occurred while tracing the function f at /usr/local/google/home/mattjj/packages/jax/froystig.py:4 for checkpoint / remat. This concrete value was not available in Python because it depends on the value of the argument x.
See https://jax.readthedocs.io/en/latest/errors.html#jax.errors.TracerBoolConversionError

Consider using the `static_argnums` parameter for `jax.remat` or `jax.checkpoint`. See the `jax.checkpoint` docstring and its example involving `static_argnums`:
https://jax.readthedocs.io/en/latest/_autosummary/jax.checkpoint.html

Notice how it says "This concrete value was not available in Python because it depends on the value of the argument x."

Also, net reduction in lines! Some cleanup that my past self promised long ago...