Allow TPU `pl.kernel` to be vmapped on nontrivial batch axes.

Also added simple batching rules for `async_copy`, which will issue one DMA call with the singular semaphore.

The primary case to support is pipelining via `emit_pipeline`, including the scalar prefetch cases.

PiperOrigin-RevId: 972088148