Supported Ops and Hardware

We currently support the following GPU kernels:

• tokamax.dot_product_attention (FlashAttention).
• tokamax.gated_linear_unit (Gated linear units (SwiGLU etc)).
• tokamax.layer_norm (Layer normalization and Root Mean Squared normalization).

And the following for both GPU and TPU:

• tokamax.ragged_dot (Mixture of Experts).