maxtext.layers package

Submodules

• maxtext.layers.attention_mla module
  • Indexer
    • Indexer.update_indexer_cache()
    • Indexer.apply_partial_rope()
    • Indexer.generate_mask()
  • mla_as_linen()
  • MLA
    • MLA.init_indexer_cache()
    • MLA.out_head_dim
    • MLA.mla_query_projection()
    • MLA.mla_get_key_value()
    • MLA.init_mla_kv_caches()
    • MLA.update_mla_kv_caches()
    • MLA.mla_kv_projection()
    • MLA.calculate_indexer_loss()
• maxtext.layers.attention_op module
  • validate_compute_axis_order()
  • apply_mask_to_logits()
  • validate_gpu_flash_attention()
  • ChunkedCausalMask
    • ChunkedCausalMask.chunk_size
  • attention_op_as_linen()
  • AttentionOp
    • AttentionOp.check_attention_inputs()
    • AttentionOp.generate_attention_mask()
    • AttentionOp.calculate_moba_gate_logic()
    • AttentionOp.generate_moba_mask_single_item()
    • AttentionOp.apply_attention()
    • AttentionOp.gpu_ragged_attention()
    • AttentionOp.tpu_ragged_attention()
    • AttentionOp.tpu_flash_attention()
    • AttentionOp.cudnn_flash_attention()
    • AttentionOp.cudnn_jax_flash_attention()
    • AttentionOp.compute_local_attention()
    • AttentionOp.is_partition_in_decode()
    • AttentionOp.apply_attention_dot()
    • AttentionOp.qk_product()
    • AttentionOp.wv_product()
    • AttentionOp.reverse_transepose()
    • AttentionOp.normalize_cudnn_attention()
    • AttentionOp.normalize_attention()
  • LoadBalancedCausalMask
    • LoadBalancedCausalMask.offset
    • LoadBalancedCausalMask.cp_size
    • LoadBalancedCausalMask.offset
• maxtext.layers.attentions module
  • L2Norm
    • L2Norm.eps
    • L2Norm.rngs
  • l2_norm_as_linen()
  • attention_as_linen()
  • Attention
    • Attention.config
    • Attention.num_query_heads
    • Attention.num_kv_heads
    • Attention.head_dim
    • Attention.max_target_length
    • Attention.mesh
    • Attention.attention_kernel
    • Attention.inputs_q_shape
    • Attention.inputs_kv_shape
    • Attention.dtype
    • Attention.weight_dtype
    • Attention.max_prefill_predict_length
    • Attention.dropout_rate
    • Attention.kernel_init
    • Attention.float32_qk_product
    • Attention.float32_logits
    • Attention.quant
    • Attention.kv_quant
    • Attention.attention_type
    • Attention.attn_logits_soft_cap
    • Attention.init_query_w()
    • Attention.query_projection()
    • Attention.init_kv_w()
    • Attention.kv_projection()
    • Attention.init_qkv_w()
    • Attention.qkv_projection()
    • Attention.out_head_dim
    • Attention.init_out_w()
    • Attention.out_projection()
    • Attention.convert_dense_general_inputs_shape()
    • Attention.init_rotary_embedding()
    • Attention.apply_rotary_embedding()
    • Attention.init_kv_caches()
    • Attention.update_kv_caches()
    • Attention.forward_serve_vllm()
• maxtext.layers.decoders module
  • DecoderLayer
    • DecoderLayer.config
    • DecoderLayer.mesh
    • DecoderLayer.model_mode
    • DecoderLayer.quant
    • DecoderLayer.name
    • DecoderLayer.parent
    • DecoderLayer.scope
  • SequentialBlockDecoderLayers
    • SequentialBlockDecoderLayers.decoder_layer
    • SequentialBlockDecoderLayers.num_decoder_layers
    • SequentialBlockDecoderLayers.config
    • SequentialBlockDecoderLayers.mesh
    • SequentialBlockDecoderLayers.quant
    • SequentialBlockDecoderLayers.model_mode
    • SequentialBlockDecoderLayers.name
    • SequentialBlockDecoderLayers.parent
    • SequentialBlockDecoderLayers.scope
  • deepstack_process()
  • Decoder
    • Decoder.config
    • Decoder.mesh
    • Decoder.quant
    • Decoder.model_mode
    • Decoder.setup()
    • Decoder.minimal_policy()
    • Decoder.get_remat_policy()
    • Decoder.get_decoder_layers()
    • Decoder.set_remat_policy()
    • Decoder.get_norm_layer()
    • Decoder.scan_decoder_layers()
    • Decoder.get_pipeline_stage_module()
    • Decoder.apply_output_head()
    • Decoder.name
    • Decoder.parent
    • Decoder.scope
• maxtext.layers.embeddings module
  • embed_as_linen()
  • Embed
    • Embed.attend()
  • attend_on_embedding()
  • rotary_embedding_as_linen()
  • RotaryEmbedding
    • RotaryEmbedding.timescale
    • RotaryEmbedding.apply_rotary()
  • llama_rotary_embedding_as_linen()
  • partial_rotary_embedding_as_linen()
  • PartialRotaryEmbedding
  • Gemma4PartialRotaryEmbedding
    • Gemma4PartialRotaryEmbedding.timescale
  • LLaMARotaryEmbedding
    • LLaMARotaryEmbedding.timescale
  • yarn_rotary_embedding_as_linen()
  • YarnRotaryEmbedding
    • YarnRotaryEmbedding.embedding_dims
    • YarnRotaryEmbedding.max_position_embeddings
    • YarnRotaryEmbedding.original_max_position_embeddings
    • YarnRotaryEmbedding.beta_fast
    • YarnRotaryEmbedding.beta_slow
    • YarnRotaryEmbedding.rope_theta
    • YarnRotaryEmbedding.rope_factor
    • YarnRotaryEmbedding.cast_as_fprop_dtype
    • YarnRotaryEmbedding.fprop_dtype
    • YarnRotaryEmbedding.rope_interleave
    • YarnRotaryEmbedding.rope_truncate
    • YarnRotaryEmbedding.rope_attention_scaling
    • YarnRotaryEmbedding.rngs
    • YarnRotaryEmbedding.freqs_cis
  • positional_embedding_as_linen()
  • PositionalEmbedding
    • PositionalEmbedding.embedding_dims
    • PositionalEmbedding.max_wavelength
    • PositionalEmbedding.cast_as_fprop_dtype
    • PositionalEmbedding.fprop_dtype
    • PositionalEmbedding.rngs
  • llama_vision_rotary_embedding_as_linen()
  • LlamaVisionRotaryEmbedding
    • LlamaVisionRotaryEmbedding.image_size
    • LlamaVisionRotaryEmbedding.patch_size
    • LlamaVisionRotaryEmbedding.hidden_size
    • LlamaVisionRotaryEmbedding.num_attention_heads
    • LlamaVisionRotaryEmbedding.rope_theta
    • LlamaVisionRotaryEmbedding.cast_as_fprop_dtype
    • LlamaVisionRotaryEmbedding.fprop_dtype
    • LlamaVisionRotaryEmbedding.rngs
    • LlamaVisionRotaryEmbedding.freqs_cis
  • Qwen3OmniMoeVisionRotaryEmbedding
    • Qwen3OmniMoeVisionRotaryEmbedding.hidden_size
    • Qwen3OmniMoeVisionRotaryEmbedding.num_attention_heads
    • Qwen3OmniMoeVisionRotaryEmbedding.spatial_merge_size
    • Qwen3OmniMoeVisionRotaryEmbedding.rope_theta
    • Qwen3OmniMoeVisionRotaryEmbedding.cast_as_fprop_dtype
    • Qwen3OmniMoeVisionRotaryEmbedding.fprop_dtype
    • Qwen3OmniMoeVisionRotaryEmbedding.rngs
    • Qwen3OmniMoeVisionRotaryEmbedding.compute_cos_sin()
  • qwen3omnimoe_vision_pos_embed_interpolate_as_linen()
  • Qwen3OmniMoeVisionPosEmbedInterpolate
    • Qwen3OmniMoeVisionPosEmbedInterpolate.num_position_embeddings
    • Qwen3OmniMoeVisionPosEmbedInterpolate.hidden_size
    • Qwen3OmniMoeVisionPosEmbedInterpolate.spatial_merge_size
    • Qwen3OmniMoeVisionPosEmbedInterpolate.dtype
    • Qwen3OmniMoeVisionPosEmbedInterpolate.cast_as_fprop_dtype
    • Qwen3OmniMoeVisionPosEmbedInterpolate.fprop_dtype
    • Qwen3OmniMoeVisionPosEmbedInterpolate.rngs
  • Qwen3OmniMoeThinkerTextRotaryEmbedding
  • qwen3_omni_mrope_embedding_as_linen()
• maxtext.layers.encoders module
  • VisionEncoder
  • AudioEncoder
  • vision_encoder_as_linen()
  • audio_encoder_as_linen()
• maxtext.layers.engram module
  • CompressedTokenizer
    • CompressedTokenizer.lookup_table
    • CompressedTokenizer.num_new_token
  • NgramHashMapping
    • NgramHashMapping.get_vocab_sizes()
  • StaticWrapper
  • MultiHeadEmbedding
  • ShortConv
  • Engram
• maxtext.layers.initializers module
  • nd_dense_init()
  • variable_to_logically_partitioned()
• maxtext.layers.learn_to_init_layer module
  • LearnToInitDecoderLayer
    • LearnToInitDecoderLayer.learn_to_init_wrapper
    • LearnToInitDecoderLayer.config
    • LearnToInitDecoderLayer.rngs
    • LearnToInitDecoderLayer.self_attention_module_name
  • LearnToInitDense
    • LearnToInitDense.C
    • LearnToInitDense.A
    • LearnToInitDense.B
    • LearnToInitDense.W
    • LearnToInitDense.bias
    • LearnToInitDense.TENSOR_A
    • LearnToInitDense.TENSOR_B
    • LearnToInitDense.TENSOR_C
    • LearnToInitDense.TENSOR_W
  • calculate_attn_weight()
  • apply_lti_model_update()
• maxtext.layers.linears module
  • normalize_axes()
  • canonicalize_tuple()
  • DenseGeneral
    • DenseGeneral.quant_dot_general
  • dense_general()
  • Dropout
  • MlpBlock
    • MlpBlock.get_norm_layer()
  • mlp_block()
• maxtext.layers.mhc module
  • get_functions()
  • sinkhorn()
  • ManifoldConstrainedHyperConnections
    • ManifoldConstrainedHyperConnections.res_mapping()
    • ManifoldConstrainedHyperConnections.mapping()
• maxtext.layers.moe module
  • get_batchsplit_init_kernel_axes()
  • random_routing()
  • calculate_load_balance_updates()
  • GateLogit
    • GateLogit.quant_dot_general
  • RoutedMoE
    • RoutedMoE.get_expert_parallelism_size()
    • RoutedMoE.get_tensor_parallelism_size()
    • RoutedMoE.get_tensor_transpose_parallelism_size()
    • RoutedMoE.get_context_autoregressive_parallelism_size()
    • RoutedMoE.should_update_load_balance()
    • RoutedMoE.get_topk()
    • RoutedMoE.deepseek_scale_weights()
    • RoutedMoE.expert_group_mask()
    • RoutedMoE.deepseek_routing()
    • RoutedMoE.apply_ffn_activation()
    • RoutedMoE.permute()
    • RoutedMoE.unpermute()
    • RoutedMoE.local_permute()
    • RoutedMoE.get_all_to_all_params()
    • RoutedMoE.transform_bias()
    • RoutedMoE.get_ragged_buffer_size()
    • RoutedMoE.sparse_matmul()
    • RoutedMoE.reshape_and_update_weights()
    • RoutedMoE.get_context_partition_and_sub_seq()
    • RoutedMoE.generate_masks_subgroup()
    • RoutedMoE.generate_masks()
    • RoutedMoE.load_balance_loss()
    • RoutedMoE.get_einsum()
    • RoutedMoE.maybe_all_gather_kernel_weight_in_expert_parallelism()
    • RoutedMoE.dense_matmul()
    • RoutedMoE.fused_moe_matmul()
    • RoutedMoE.retrieve_quantized_weight()
  • RoutedAndSharedMoE
    • RoutedAndSharedMoE.routed_moe
  • get_gate_logit()
  • get_routed_moe()
  • get_routed_and_shared_moe()
• maxtext.layers.multi_token_prediction module
  • mtp_losses
  • mtp_acceptance
  • roll_and_mask()
  • MultiTokenPredictionLayer
    • MultiTokenPredictionLayer.embedding_norm
    • MultiTokenPredictionLayer.hidden_state_norm
    • MultiTokenPredictionLayer.projection_layer
    • MultiTokenPredictionLayer.transformer_layer
  • MultiTokenPredictionBlock
  • calculate_mtp_loss()
  • calculate_mtp_acceptance_rate()
  • multi_token_prediction_block_as_linen()
• maxtext.layers.nnx_decoders module
  • NNXDecoderLayer
  • deepstack_process()
  • NNXDecoder
    • NNXDecoder.get_decoder_layers()
    • NNXDecoder.minimal_policy()
    • NNXDecoder.get_remat_policy()
    • NNXDecoder.get_norm_layer()
    • NNXDecoder.apply_output_head()
  • decoder_as_linen()
• maxtext.layers.nnx_wrappers module
  • is_vanilla_variable()
  • to_linen_var()
  • get_col_name()
  • to_nnx_var()
  • linen_vars_to_nnx_attrs()
  • nnx_attrs_to_linen_vars()
  • lazy_init()
  • current_linen_module()
  • ToNNX
    • ToNNX.lazy_init()
  • linen_rngs_dict()
  • ToLinen
    • ToLinen.nnx_class
    • ToLinen.args
    • ToLinen.kwargs
    • ToLinen.skip_rng
    • ToLinen.metadata_fn()
    • ToLinen.name
    • ToLinen.parent
    • ToLinen.scope
  • to_linen()
  • to_linen_class()
• maxtext.layers.normalizations module
  • RMSNorm
  • GlobalRMSNorm
  • Qwen3NextRMSNorm()
  • Qwen3NextRMSNormGated
    • Qwen3NextRMSNormGated.num_features
    • Qwen3NextRMSNormGated.eps
    • Qwen3NextRMSNormGated.dtype
    • Qwen3NextRMSNormGated.weight_dtype
  • rms_norm()
  • l2norm()
• maxtext.layers.pipeline module
  • PipelineBase
    • PipelineBase.config
    • PipelineBase.layers
    • PipelineBase.mesh
    • PipelineBase.remat_policy
    • PipelineBase.setup()
    • PipelineBase.need_circ_storage()
    • PipelineBase.iterations_to_complete_first_microbatch_one_repeat()
    • PipelineBase.iterations_to_complete_first_microbatch()
    • PipelineBase.get_iteration_inputs()
    • PipelineBase.get_microbatch_and_repeat_ids()
    • PipelineBase.get_pipeline_remat_policy()
    • PipelineBase.get_weight_sharding()
    • PipelineBase.get_vmap_func_for_init()
    • PipelineBase.get_main_vmap_func_for_iterations()
    • PipelineBase.name
    • PipelineBase.parent
    • PipelineBase.scope
  • Pipeline
    • Pipeline.init_states()
    • Pipeline.shard_dim_by_stages()
    • Pipeline.vmap_parallel_gather()
    • Pipeline.vmap_gather()
    • Pipeline.get_new_loop_state()
    • Pipeline.permute_output_micro_per_stage_dim()
    • Pipeline.get_current_stage_weights()
    • Pipeline.get_current_repeat_from_stages()
    • Pipeline.run_one_iteration()
    • Pipeline.get_logical_spec_repeats_removed()
    • Pipeline.all_gather_over_fsdp()
    • Pipeline.name
    • Pipeline.parent
    • Pipeline.scope
  • CircularPipeline
    • CircularPipeline.init_states()
    • CircularPipeline.gather_weights_across_stages_vmap()
    • CircularPipeline.gather_microbatch_inputs_vmap()
    • CircularPipeline.advance_circular_buffers()
    • CircularPipeline.realign_output_microbatches()
    • CircularPipeline.fetch_active_stage_weights()
    • CircularPipeline.get_current_weights_from_bsw()
    • CircularPipeline.from_all_variables_to_repeat_weights()
    • CircularPipeline.from_repeat_weights_to_bsw()
    • CircularPipeline.weight_prefetching()
    • CircularPipeline.run_one_iteration()
    • CircularPipeline.name
    • CircularPipeline.parent
    • CircularPipeline.scope
  • create_pipeline()
• maxtext.layers.pipeline_deprecated module
  • Pipeline
    • Pipeline.config
    • Pipeline.layers
    • Pipeline.mesh
    • Pipeline.remat_policy
    • Pipeline.setup()
    • Pipeline.need_circ_storage()
    • Pipeline.iterations_to_complete_first_microbatch_one_repeat()
    • Pipeline.iterations_to_complete_first_microbatch()
    • Pipeline.init_states()
    • Pipeline.get_iteration_inputs()
    • Pipeline.shard_dim_by_stages()
    • Pipeline.get_microbatch_and_repeat_ids()
    • Pipeline.vmap_parallel_gather()
    • Pipeline.vmap_gather()
    • Pipeline.get_new_loop_state()
    • Pipeline.permute_output_micro_per_stage_dim()
    • Pipeline.get_current_stage_weights()
    • Pipeline.get_current_repeat_from_stages()
    • Pipeline.get_vmap_func_for_init()
    • Pipeline.get_main_vmap_func_for_iterations()
    • Pipeline.run_one_iteration()
    • Pipeline.get_pipeline_remat_policy()
    • Pipeline.get_weight_sharding()
    • Pipeline.get_logical_spec_repeats_removed()
    • Pipeline.all_gather_over_fsdp()
    • Pipeline.name
    • Pipeline.parent
    • Pipeline.scope
• maxtext.layers.quantizations module
  • Quantization
    • Quantization.dot_general_cls()
    • Quantization.einsum()
  • AqtQuantization
    • AqtQuantization.quant_dg
    • AqtQuantization.quant_mode
    • AqtQuantization.replicate_scale
    • AqtQuantization.dot_general_cls()
    • AqtQuantization.einsum()
  • QwixQuantization
    • QwixQuantization.quant_mode
    • QwixQuantization.act_calibration_method
    • QwixQuantization.weight_calibration_method
    • QwixQuantization.bwd_calibration_method
    • QwixQuantization.dot_general_cls()
    • QwixQuantization.einsum()
  • QwixDotGeneral
    • QwixDotGeneral.config
    • QwixDotGeneral.name
    • QwixDotGeneral.parent
    • QwixDotGeneral.scope
  • QwixEinsum
    • QwixEinsum.config
    • QwixEinsum.name
    • QwixEinsum.parent
    • QwixEinsum.scope
  • Fp8Quantization
    • Fp8Quantization.quant_mode
    • Fp8Quantization.dot_general_cls()
    • Fp8Quantization.einsum()
  • Fp8Einsum
    • Fp8Einsum.amax_history_length
    • Fp8Einsum.e4m3_dtype
    • Fp8Einsum.e5m2_dtype
    • Fp8Einsum.dtype
    • Fp8Einsum.setup()
    • Fp8Einsum.name
    • Fp8Einsum.parent
    • Fp8Einsum.scope
  • NANOOFp8Quantization
    • NANOOFp8Quantization.quant_mode
    • NANOOFp8Quantization.dot_general_cls()
  • ConstantBoundConfig
    • ConstantBoundConfig.fwd_lhs_bound
    • ConstantBoundConfig.fwd_rhs_bound
    • ConstantBoundConfig.dlhs_lhs_bound
    • ConstantBoundConfig.dlhs_rhs_bound
    • ConstantBoundConfig.drhs_lhs_bound
    • ConstantBoundConfig.drhs_rhs_bound
  • PerTensorScales
    • PerTensorScales.fwd_lhs
    • PerTensorScales.fwd_rhs
    • PerTensorScales.dlhs_lhs
    • PerTensorScales.dlhs_rhs
    • PerTensorScales.drhs_lhs
    • PerTensorScales.drhs_rhs
  • in_convert_mode()
  • in_serve_mode()
  • get_quant_mode()
  • configure_quantization()
  • match_aqt_and_unquantized_param()
  • remove_quantized_params()
  • configure_kv_quant()
  • NvidaFp8Provider
    • NvidaFp8Provider.dot_general()
    • NvidaFp8Provider.einsum()
  • NANOOFp8Provider
    • NANOOFp8Provider.dot_general()
  • get_fp8_full_qwix_rule_w_sparsity()
  • get_quantization_rule()
  • get_qt_provider()
  • maybe_quantize_model()
  • manual_quantize()
  • TransformerEngineQuantization
    • TransformerEngineQuantization.get_block_size()
    • TransformerEngineQuantization.dot_general_cls()
    • TransformerEngineQuantization.einsum()
• maxtext.layers.train_state_nnx module
  • TrainStateNNX
    • TrainStateNNX.apply_gradients()