jax

323353eb - implement higher-rank cases of lax_numpy.dot in terms of lax.dot_general, avoiding a reshape