maxtext.models.gemma4 module

maxtext.models.gemma4 module#

Specialized layers for Gemma 4.

maxtext.models.gemma4.get_attention_type(layer_id)[source]#
class maxtext.models.gemma4.Gemma4MoE(*args, **kwargs)[source]#

Bases: Module

Gemma4 specific MoE block containing layer norms and a generic MoE block.

Parameters:

  • args (Any)

  • kwargs (Any)

Return type:

Any

class maxtext.models.gemma4.Gemma4DecoderLayer(*args, **kwargs)[source]#

Bases: Module

Transformer decoder layer for Gemma4.

Parameters:

  • args (Any)

  • kwargs (Any)

Return type:

Any

class maxtext.models.gemma4.Gemma4ScannableBlock(*args, **kwargs)[source]#

Bases: Module

A repeatable block of Gemma4 decoder layers.

Parameters:

  • args (Any)

  • kwargs (Any)

Return type:

Any