jax

99944d12 - Fix lax.broadcast_shapes returning numpy ints in shape tuple (#2471)

Fix lax.broadcast_shapes returning numpy ints in shape tuple (#2471) * Fix lax.broadcast_shapes returning numpy ints in shape tuple * Use _canonicalize_dimension and add test