mlp_gpt_jax
