[Feature] Nested calls to accelerate should behave like nested jax.jit functions.

[josh146]

For example, consider the following:

@qjit
def f(x):
    return accelerate(accelerate(jnp.sin))(x)

This works fine:

>>> f(0.43)
Array(0.4168708, dtype=float64)

However, differentiating it fails:

>>> qjit(grad(f))(0.43)
ValueError: Function vjp_wrapper must be jax.jit-able.But failed with error message Differentiation rule for 'python_callback' not implemented.

We should treat accelerate the same way that JAX treats jit; that is, a nested pjit context in the jaxpr. This will allow for arbitrary nesting of arbitrary jax-jitable code, without breaking gradients.