Skip to main content
Log in

Python module

naive_transformer

NaiveTransformer

class max.pipelines.nn.transformer.naive_transformer.NaiveTransformer(dim: int, n_heads: int, layers: list[max.pipelines.nn.transformer.naive_transformer.NaiveTransformerBlock], norm: Layer, output: Linear | LinearV2, theta: float, embedding: Embedding | EmbeddingV2, output_type: DType | None = None)

Max-Graph only model consisting of NaiveTransformerBlock layers.

NaiveTransformerBlock

class max.pipelines.nn.transformer.naive_transformer.NaiveTransformerBlock(attention: NaiveAttentionWithRope, mlp: Layer, attention_norm: Layer, mlp_norm: Layer)

Max-Graph Only Stack of Attention, FeedForward, and RMSNorm layers.