jax

ed0d8c02 - 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.)