tweak lax.py shape broadcasting logic (#4217)
tweak lax.py shape broadcasting logic (#4217)

This new implementation is faster, and works for polymorphic shapes without weird tricks. (This new implementation is faster even if we remove the weird tricks for polymorphism.)
google/jax,tensorflow/probability,tensorflow/probability,google/jax,google/jax,google/jax