jax

8fc1332e - Avoid broadcasting for batch dimensions in jnp.matmul.

Avoid broadcasting for batch dimensions in jnp.matmul. Instead, squeeze size 1 dimensions out of the matmul input, and transpose any non-batch non-contracting dimensions into the correct location.