maxtext.models.gemma2 module

Specialised layers for Gemma2.

class maxtext.models.gemma2.Gemma2DecoderLayer(*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