fake_pmap_and_jit has a confusing interface

I spend quite some time figuring out why code in a large codebase was so slow, only to find out that `jit` was disabled throughout the entire project. This was because the `main` function was called as follows:

```py
with chex.fake_pmap_and_jit(FLAGS.debug):
main()
```

While on first sight it appears as if this indeed disables both pmap and jit if flag `debug` is set, this in fact only disables `pmap` and _always disables jit_!

The reason is that `fake_pmap_and_jit` take two positional arguments that disable respectively `pmap` and `jit`, and they are both `True` by default. The names of these arguments are somewhat cryptic to me as well: `enable_pmap_patching` and `enable_jit_patching`, which actually *disable* these JAX transformations.

Given these observations, I think the situation would improve if the signature would be:

```py
def fake_pmap_and_jit(*, disable_pmap: bool = True, disable_jit: bool = True)
```

Then my code above would then look like this:

```py
with chex.fake_pmap_and_jit(disable_pmap=FLAGS.debug):
main()
```

Which shows clearly we are not setting `disable_jit`, so we would rewrite this to:

```py
with chex.fake_pmap_and_jit(disable_pmap=FLAGS.debug, disable_jit=FLAGS.debug):
main()
```