maxtext.models package

Submodules

• maxtext.models.deepseek module
  • DeepSeekGenericLayer
    • DeepSeekGenericLayer.mlp_op()
    • DeepSeekGenericLayer.with_logical_constraint()
    • DeepSeekGenericLayer.dropout_op()
    • DeepSeekGenericLayer.pre_attention_norm_op()
    • DeepSeekGenericLayer.post_attention_norm_op()
    • DeepSeekGenericLayer.attention_op()
    • DeepSeekGenericLayer.logical_axis_names
    • DeepSeekGenericLayer.mlp_logical_axis_names
    • DeepSeekGenericLayer.post_process()
    • DeepSeekGenericLayer.self_attention_with_norm_op()
    • DeepSeekGenericLayer.engram_op()
  • DeepSeekDenseLayer
    • DeepSeekDenseLayer.mlp_op()
  • DeepSeekMoELayer
    • DeepSeekMoELayer.mlp_op()
• maxtext.models.deepseek_batchsplit module
  • scheduling_group()
  • fetch_weights()
  • split()
  • merge()
  • extract_layer_weights()
  • insert_layer_ws_grad()
  • gather_weights()
  • reduce_scatter_ws_grad()
  • all_reduce_ws_grad_dcn()
  • init_splash_kernel()
  • tpu_flash_attention()
  • tpu_flash_attention_bwd()
  • scan_batch_split_layers()
  • batch_split_schedule()
  • batch_split_schedule_bwd()
  • staggered_call()
  • dot()
  • mla_with_norms()
  • mla_with_norms_remat()
  • mla_with_norms_bwd()
  • mla()
  • mla_remat()
  • mla_bwd()
  • query_projection()
  • kv_projection()
  • get_key_value()
  • rms_norm()
  • initialize_yarn_mask()
  • initialize_yarn_freqs()
  • yarn()
  • shared_expert_and_route()
  • shared_expert()
  • expert_group_mask()
  • expert_indices_and_weights()
  • expert_selection()
  • route()
  • unroute()
  • route_impl_fwd()
  • route_impl_bwd()
  • unroute_impl_fwd()
  • unroute_impl_bwd()
  • gmm()
  • compute_gating()
  • compute_linear()
  • route_compute_unroute()
  • unroute_ubatch_shard_mapped()
  • unroute_ubatch_fn()
  • unroute_ubatch_remat_and_bwd_shard_mapped()
  • unroute_ubatch_fn_remat()
  • unroute_ubatch_fn_bwd()
  • sum_grads()
  • route_compute_unroute_bwd()
  • moe()
  • moe_bwd()
• maxtext.models.deepseek_batchsplit_fp8 module
  • fetch_weights()
  • split()
  • merge()
  • gather_weights()
  • scan_batch_split_layers()
  • batch_split_schedule()
  • staggered_call()
  • with_data_parallel_constraint()
  • dot()
  • mla_with_norms()
  • mla()
  • query_projection()
  • kv_projection()
  • get_key_value()
  • rms_norm()
  • yarn()
  • moe()
  • expert_indices_and_weights()
  • expert_selection()
  • route()
  • unroute()
  • compute()
  • route_compute_unroute()
  • process_activations()
• maxtext.models.gemma module
  • GemmaDecoderLayer
    • GemmaDecoderLayer.config
    • GemmaDecoderLayer.mesh
    • GemmaDecoderLayer.model_mode
    • GemmaDecoderLayer.quant
• maxtext.models.gemma2 module
  • Gemma2DecoderLayer
    • Gemma2DecoderLayer.config
    • Gemma2DecoderLayer.mesh
    • Gemma2DecoderLayer.model_mode
    • Gemma2DecoderLayer.quant
• maxtext.models.gemma3 module
  • get_attention_type()
  • get_query_pre_attn_scalar()
  • Gemma3DecoderLayer
  • Gemma3ScannableBlock
  • MlpBlockViT
  • Encoder1DBlock
  • Encoder
  • Einsum
  • VisionEmbedder
  • visionembedder_as_linen()
  • VisionExit
    • VisionExit.output_length
  • vision_exit_as_linen()
  • Gemma3VisionEncoderLayer
  • gemma3visionencoder_as_linen()
• maxtext.models.gemma4 module
  • get_attention_type()
  • Gemma4MoE
  • Gemma4DecoderLayer
  • Gemma4ScannableBlock
• maxtext.models.gemma4_vision module
  • factorized_posemb()
  • patchify()
  • VisionEntry
  • apply_multidimensional_rope()
  • avg_pool_by_positions()
  • VisionExit
  • Gemma4VisionRotaryEmbedding
  • Gemma4Attention
    • Gemma4Attention.init_rotary_embedding()
  • Gemma4EncoderBlock
  • Gemma4VisionEncoderLayer
  • Gemma4VisionProjector
  • gemma4_vision_encoder_as_linen()
• maxtext.models.gpt3 module
  • Gpt3LayerNorm
  • gpt3_layer_norm()
  • Gpt3MultiHeadAttention
    • Gpt3MultiHeadAttention.num_heads
    • Gpt3MultiHeadAttention.head_dim
    • Gpt3MultiHeadAttention.max_target_length
    • Gpt3MultiHeadAttention.max_prefill_predict_length
    • Gpt3MultiHeadAttention.mesh
    • Gpt3MultiHeadAttention.dtype
    • Gpt3MultiHeadAttention.dropout_rate
    • Gpt3MultiHeadAttention.kernel_init
    • Gpt3MultiHeadAttention.float32_qk_product
    • Gpt3MultiHeadAttention.float32_logits
    • Gpt3MultiHeadAttention.fused_qkv
    • Gpt3MultiHeadAttention.quant
    • Gpt3MultiHeadAttention.use_bias
    • Gpt3MultiHeadAttention.create_projection_layer()
    • Gpt3MultiHeadAttention.qkv_projection()
    • Gpt3MultiHeadAttention.projection()
  • Gpt3DecoderLayer
• maxtext.models.gpt_oss module
  • get_attention_type()
  • GptOssDecoderLayer
  • GptOssScannableBlock
    • GptOssScannableBlock.config
    • GptOssScannableBlock.mesh
    • GptOssScannableBlock.num_of_layers
    • GptOssScannableBlock.quant
• maxtext.models.llama2 module
  • LlamaDecoderLayer
  • LlamaLTIDecoderLayer
• maxtext.models.llama4 module
  • Llama4UnfoldConvolution
    • Llama4UnfoldConvolution.config
  • pixel_shuffle()
  • Llama4VisionMLP
    • Llama4VisionMLP.config
  • Llama4VisionMLP2
    • Llama4VisionMLP2.config
  • Llama4VisionPixelShuffleMLP
    • Llama4VisionPixelShuffleMLP.config
  • Llama4MultiModalProjector
    • Llama4MultiModalProjector.config
  • llama4multimodalprojector_as_linen()
  • determine_is_nope_layer()
  • determine_is_moe_layer()
  • Llama4DecoderLayer
    • Llama4DecoderLayer.moe_block
  • Llama4ScannableBlock
  • Llama4VisionEncoderLayer
  • Llama4VisionEncoder
    • Llama4VisionEncoder.config
    • Llama4VisionEncoder.mesh
  • Llama4VisionModel
    • Llama4VisionModel.config
    • Llama4VisionModel.mesh
  • llama4visionmodel_as_linen()
• maxtext.models.mistral module
  • MistralDecoderLayer
• maxtext.models.mixtral module
  • MixtralDecoderLayer
• maxtext.models.models module
  • TransformerLinenPure
    • TransformerLinenPure.config
    • TransformerLinenPure.mesh
    • TransformerLinenPure.quant
    • TransformerLinenPure.model_mode
    • TransformerLinenPure.init()
    • TransformerLinenPure.apply()
    • TransformerLinenPure.setup()
    • TransformerLinenPure.logits_from_hidden_states()
    • TransformerLinenPure.name
    • TransformerLinenPure.parent
    • TransformerLinenPure.scope
  • transformer_as_linen()
  • TransformerLinen
    • TransformerLinen.init()
    • TransformerLinen.apply()
    • TransformerLinen.name
    • TransformerLinen.parent
    • TransformerLinen.scope
  • Transformer
    • Transformer.no_op()
    • Transformer.init_cache()
• maxtext.models.olmo3 module
  • get_attention_type()
  • Olmo3DecoderLayer
  • Olmo3ScannableBlock
    • Olmo3ScannableBlock.config
    • Olmo3ScannableBlock.mesh
    • Olmo3ScannableBlock.num_of_layers
    • Olmo3ScannableBlock.quant
• maxtext.models.qwen2 module
  • AttentionWithNorm
    • AttentionWithNorm.apply_attention_with_norm()
  • Qwen2DecoderLayer
• maxtext.models.qwen3 module
  • naive_jax_chunk_gated_delta_rule()
  • jax_chunk_gated_delta_rule()
  • Qwen3NextGatedDeltaNet
  • Qwen3NextFullAttention
    • Qwen3NextFullAttention.config
    • Qwen3NextFullAttention.mesh
    • Qwen3NextFullAttention.model_mode
    • Qwen3NextFullAttention.layer_idx
    • Qwen3NextFullAttention.quant
    • Qwen3NextFullAttention.attention
  • Qwen3NextSparseMoeBlock
    • Qwen3NextSparseMoeBlock.config
    • Qwen3NextSparseMoeBlock.mesh
    • Qwen3NextSparseMoeBlock.quant
  • Qwen3NextScannableBlock
    • Qwen3NextScannableBlock.config
    • Qwen3NextScannableBlock.mesh
    • Qwen3NextScannableBlock.model_mode
    • Qwen3NextScannableBlock.quant
  • Qwen3NextDecoderLayer
    • Qwen3NextDecoderLayer.config
    • Qwen3NextDecoderLayer.mesh
    • Qwen3NextDecoderLayer.model_mode
    • Qwen3NextDecoderLayer.layer_idx
    • Qwen3NextDecoderLayer.quant
  • AttentionWithNorm
    • AttentionWithNorm.apply_attention_with_norm()
  • Qwen3DecoderLayer
  • Qwen3MoeDecoderLayer
  • Qwen3OmniMoeVisionPatchMerger
    • Qwen3OmniMoeVisionPatchMerger.config
    • Qwen3OmniMoeVisionPatchMerger.hidden_size
    • Qwen3OmniMoeVisionPatchMerger.use_postshuffle_norm
    • Qwen3OmniMoeVisionPatchMerger.dtype
    • Qwen3OmniMoeVisionPatchMerger.weight_dtype
    • Qwen3OmniMoeVisionPatchMerger.kernel_init
    • Qwen3OmniMoeVisionPatchMerger.rngs
    • Qwen3OmniMoeVisionPatchMerger.ln_q
    • Qwen3OmniMoeVisionPatchMerger.mlp_0
    • Qwen3OmniMoeVisionPatchMerger.mlp_2
  • Qwen3OmniMoeVisionMLP
    • Qwen3OmniMoeVisionMLP.config
    • Qwen3OmniMoeVisionMLP.hidden_size
    • Qwen3OmniMoeVisionMLP.intermediate_size
    • Qwen3OmniMoeVisionMLP.dtype
    • Qwen3OmniMoeVisionMLP.weight_dtype
    • Qwen3OmniMoeVisionMLP.kernel_init
    • Qwen3OmniMoeVisionMLP.rngs
    • Qwen3OmniMoeVisionMLP.linear_fc1
    • Qwen3OmniMoeVisionMLP.linear_fc2
  • Qwen3OmniMoeVisionPatchEmbed
    • Qwen3OmniMoeVisionPatchEmbed.config
    • Qwen3OmniMoeVisionPatchEmbed.patch_size
    • Qwen3OmniMoeVisionPatchEmbed.temporal_patch_size
    • Qwen3OmniMoeVisionPatchEmbed.in_channels
    • Qwen3OmniMoeVisionPatchEmbed.embed_dim
    • Qwen3OmniMoeVisionPatchEmbed.dtype
    • Qwen3OmniMoeVisionPatchEmbed.weight_dtype
    • Qwen3OmniMoeVisionPatchEmbed.rngs
    • Qwen3OmniMoeVisionPatchEmbed.proj
  • Qwen3OmniMoeVisionAttention
    • Qwen3OmniMoeVisionAttention.config
    • Qwen3OmniMoeVisionAttention.attn
  • Qwen3OmniMoeVisionBlock
    • Qwen3OmniMoeVisionBlock.config
    • Qwen3OmniMoeVisionBlock.ln1
    • Qwen3OmniMoeVisionBlock.ln2
    • Qwen3OmniMoeVisionBlock.attn
    • Qwen3OmniMoeVisionBlock.mlp
    • Qwen3OmniMoeVisionBlock.mlp_out
  • Qwen3OmniMoeVisionEncoder
    • Qwen3OmniMoeVisionEncoder.config
    • Qwen3OmniMoeVisionEncoder.patch_embed
    • Qwen3OmniMoeVisionEncoder.pos_embed_interpolate
    • Qwen3OmniMoeVisionEncoder.blocks
    • Qwen3OmniMoeVisionEncoder.merger_list
    • Qwen3OmniMoeVisionEncoder.spatial_merge_size
    • Qwen3OmniMoeVisionEncoder.deep_idx
  • Qwen3OmniMoeVisionProjector
    • Qwen3OmniMoeVisionProjector.config
    • Qwen3OmniMoeVisionProjector.merger
  • qwen3omni_visionencoder_as_linen()
  • qwen3omni_visionprojector_as_linen()
  • Qwen3OmniAudioEncoderLayer
  • Qwen3OmniAudioEncoder
    • Qwen3OmniAudioEncoder.config
    • Qwen3OmniAudioEncoder.mesh
  • Qwen3OmniAudioProjector
  • qwen3omni_audioencoder_as_linen()
  • qwen3omni_audioprojector_as_linen()
• maxtext.models.qwen3_5 module
  • Qwen3_5GatedDeltaNet
  • Qwen3_5FullAttention
  • Qwen3_5SparseMoEBlock
  • Qwen3_5ScannableBlock
  • Qwen3_5DecoderLayer
    • Qwen3_5DecoderLayer.config
    • Qwen3_5DecoderLayer.mesh
    • Qwen3_5DecoderLayer.model_mode
    • Qwen3_5DecoderLayer.layer_idx
    • Qwen3_5DecoderLayer.quant
• maxtext.models.qwen3_custom module
  • Qwen3CustomAttention
    • Qwen3CustomAttention.init_out_w()
  • Qwen3CustomMoeDecoderLayer
    • Qwen3CustomMoeDecoderLayer.apply_attention_with_norm()
• maxtext.models.simple_layer module
  • SimpleDecoderLayer
  • SimpleMlpDecoderLayer