Why lax.dot_general instead of jnp.dot

[patil-suraj]

The nn.Dense layer uses lax.dot_general for the dot product. Why is it preferred over jnp.dot ?

[jheek]

Actually I think we could just switch to dot now before we used dot_general because dot did not have the precision argument but that's no longer the case.

[jheek]

Feel free to open an issue or PR

[8bitmp3]

@jheek @avital @marcvanzee @levskaya Would it be better if one of the core devs implements the changes to speed things up and avoid breaking something for the next release? jnp.dot seems like a natural fit and less confusing. Thanks! 👍