maxtext.models.llama2 module

maxtext.models.llama2 module#

Transformer model definition.

class maxtext.models.llama2.LlamaDecoderLayer(*args, **kwargs)[source]#

Bases: Module

Transformer decoder layer that attends to the encoder.

Parameters:

  • args (Any)

  • kwargs (Any)

Return type:

Any

class maxtext.models.llama2.LlamaLTIDecoderLayer(*args, **kwargs)[source]#

Bases: LearnToInitDecoderLayer

A Type-bounded version of Llama-specific LearnToInitDecoderLayer. Temporal LTI wrapper before it is generalized for other models.

Parameters:

  • args (Any)

  • kwargs (Any)

Return type:

Any