maxtext.models.gemma module

Specialised layers for Gemma.

class maxtext.models.gemma.GemmaDecoderLayer(*args, **kwargs)

Bases: Module

Transformer decoder layer that attends to the encoder.

Parameters:

• args (Any)

• kwargs (Any)

Return type:

Any

config: Any

mesh: Mesh

model_mode: str

quant: None | AqtQuantization = None