rename jax/batched_stuff to jax/chunked_stuff

Created by: PhilipVinc

As agreed in the other PR. I split this out to have a less messy PR, and merge it immediately so @inailuig can rebase SR on top of this change