It appears the JAX version of the codebase already supports a fused GEGLU implementation (as introduced in #602). Would it be possible to port over the optimized codepath for SwiGLU, to support a (swish, linear) use case (as the current (gelu, linear) case?)
It appears the JAX version of the codebase already supports a fused GEGLU implementation (as introduced in #602). Would it be possible to port over the optimized codepath for SwiGLU, to support a
(swish, linear)use case (as the current(gelu, linear)case?)