improve broadcast handling in batching.py
improve broadcast handling in batching.py

We can avoid a circular import just by `import jax`!

google/jax,google/jax,google/jax,google/jax,tensorflow/probability,tensorflow/probability