jax

b6d4fe53 - Define lax.ragged_dot_general and express lax.ragged_dot in terms of it.