search

AI-Hypercomputer / maxtext

A simple, performant and scalable Jax LLM!
Apache License 2.0
1.47k stars 275 forks source link

Add simple MLP decoder block #858

Closed gobbleturk closed 3 weeks ago

gobbleturk commented 3 weeks ago

Add a simple [embed,mlp] @ [mlp,embed] decoder block for debugging and testing purposes