Switch eqx.nn.Linear.__call__ to x @ weight.T to eliminate vmap+grad permutes (gated on GPU benchmark)#1235
Open
mvanhorn wants to merge 2 commits into