maxtext package

MaxText is a high performance, highly scalable, open-source LLM written in pure Python/Jax and targeting Google Cloud TPUs and GPUs for training and inference. MaxText achieves high MFUs and scales from single host to very large clusters while staying simple and “optimization-free” thanks to the power of Jax and the XLA compiler.

Subpackages

• maxtext.checkpoint_conversion package
  • Subpackages
    • maxtext.checkpoint_conversion.utils package
      • Submodules
  • Submodules
    • maxtext.checkpoint_conversion.compare_hf_ckpt module
      • load_safetensors_generic()
      • get_hf_model_state_dict()
      • verify_dictionaries()
      • main()
    • maxtext.checkpoint_conversion.load_and_quantize_checkpoint module
    • maxtext.checkpoint_conversion.reshard_checkpoint module
    • maxtext.checkpoint_conversion.to_huggingface module
    • maxtext.checkpoint_conversion.to_maxtext module
      • LazyHFLoader
      • LazyTensor
      • LazyTensorHandler
      • get_maxtext_model_info()
      • main()
• maxtext.configs package
  • Submodules
    • maxtext.configs.pyconfig module
      • initialize()
      • initialize_pydantic()
    • maxtext.configs.pyconfig_deprecated module
      • yaml_key_to_env_key()
      • string_to_bool()
      • validate_compute_axis_order()
      • validate_shard_mode()
      • validate_kv_quant_axis()
      • validate_attention_kernel()
      • validate_attention_type()
      • validate_moba_attention()
      • validate_attention_window_params()
      • validate_profiler_type()
      • validate_periodic_profiler()
      • validate_model_call_mode()
      • validate_prefill_and_target_lengths()
      • validate_rope_type()
      • validate_expert_shard_attention_option()
      • validate_vocab_tiling()
      • validate_rampup_batch_size()
      • validate_context_parallel_strategy_ring()
      • validate_keys()
      • validate_tokenizer()
      • validate_constant_bound()
      • validate_quantization_methods()
      • validate_tokamax_usage()
      • validate_data_input()
      • validate_llama4_config()
      • validate_model_name()
      • validate_multimodal_model_name()
      • validate_no_keys_overwritten_twice()
      • validate_and_assign_remat_tensors()
      • resolve_config_path()
      • create_parallelisms_list()
      • set_mu_dtype()
      • validate_and_set_hlo_dump_defaults()
      • validate_multiple_slices()
      • set_and_validate_pipeline_config()
      • validate_deepseek_moe()
      • validate_mlp_dim()
      • validate_gpt_oss_moe()
      • validate_sparse_matmul_parallelism()
      • validate_ring_of_experts_parallelism()
      • validate_shard_expert_on_fsdp()
      • validate_ragged_dot()
      • validate_optimizer_sharding_over_data()
      • create_new_logical_axis_rules()
      • update_model_keys()
      • validate_and_update_keys()
      • get_individual_scales()
      • calculate_global_batch_sizes()
      • calculate_rampup_samples_and_steps()
      • get_num_target_devices()
      • get_quantization_local_shard_count()
      • get_context_parallel_size()
      • using_pipeline_parallelism()
      • using_tensor_parallelism()
      • using_sequence_parallelism()
      • using_expert_parallelism()
      • using_fsdp_and_transpose_parallelism()
      • HyperParameters
      • initialize()
    • maxtext.configs.types module
      • XProfTPUPowerTraceMode
      • DType
      • MatmulPrecision
      • QuantizationType
      • KvQuantAxis
      • RematPolicy
      • RematLocation
      • OptimizerType
      • LearningRateScheduleType
      • WsdDecayStyle
      • RopeType
      • TokenizerType
      • DatasetType
      • SamplingStrategy
      • ProfilerType
      • RunInfo
      • Checkpointing
      • OrbaxStorage
      • EmergencyCheckpointing
      • DataTypes
      • Quantization
      • ModelArchitecture
      • MTP
      • Logits
      • Attention
      • MoBa
      • MlaAttention
      • AttentionIndexer
      • Llama4Attention
      • SplashAttention
      • PagedAttention
      • MoEGeneral
      • MoEKernels
      • DeepSeekMoE
      • Qwen3Next
      • HardwareAndMesh
      • LayoutAndSharding
      • DcnParallelism
      • IciParallelism
      • PipelineParallelism
      • RematAndOffload
      • Tokenizer
      • DatasetGeneral
      • TfdsDataset
      • HfDataset
      • GrainDataset
      • OlmoGrainDataset
      • FineTuning
      • Distillation
      • TrainingLoop
      • ManifoldConstrainedHyperConnections
      • DilocoParams
      • Optimizer
      • AdamW
      • Muon
      • PositionalEmbedding
      • Rope
      • YarnRope
      • InferenceGeneral
      • Decoding
      • InferenceLayout
      • InferenceServer
      • InferenceBenchmark
      • PrefixCaching
      • AOT
      • DevelopmentAndDebugging
      • Profiling
      • HloDump
      • StackTrace
      • Metrics
      • ManagedMLDiagnostics
      • Goodput
      • ElasticTraining
      • GcpMonitoring
      • Tensorboard
      • MultimodalGeneral
      • VisionTower
      • VisionProjector
      • AudioEncoder
      • Debug
      • RLHardware
      • VLLM
      • RL
      • RLDataset
      • RLEvaluation
      • Reward
      • SpecialTokens
      • Engram
      • DerivedValues
      • get_individual_scales()
      • MaxTextConfig
• maxtext.experimental package
  • Subpackages
    • maxtext.experimental.agent package
      • Subpackages
    • maxtext.experimental.rl package
      • Submodules
• maxtext.input_pipeline package
  • Subpackages
    • maxtext.input_pipeline.packing package
      • Submodules
    • maxtext.input_pipeline.protos package
      • Submodules
  • Submodules
    • maxtext.input_pipeline.data_processing_utils module
      • parse_and_keep_features()
      • get_tokenizer_and_pad_id()
      • validate_and_configure_sft_columns()
      • get_local_batch_size()
      • format_and_batch()
      • apply_multiprocessing_and_prefetch()
    • maxtext.input_pipeline.distillation_data_processing module
      • InputRequest
      • map_to_prompt_completion()
      • extract_content()
      • process_dataset()
      • load_dataset()
      • filter_dataset()
    • maxtext.input_pipeline.grain_data_processing module
      • find_data_files()
      • get_datasets()
      • pretrain_preprocessing_pipeline()
      • dpo_preprocessing_pipeline()
      • sft_preprocessing_pipeline()
      • make_grain_train_iterator()
      • make_grain_eval_iterator()
    • maxtext.input_pipeline.grain_tokenizer module
      • TokenizerTransformBase
      • TokenizeAndTrim
      • TokenizeAndChunk
    • maxtext.input_pipeline.hf_data_processing module
      • vision_sft_preprocessing_pipeline()
      • preprocessing_pipeline()
      • make_hf_train_iterator()
      • make_hf_eval_iterator()
    • maxtext.input_pipeline.input_pipeline_interface module
      • get_process_loading_real_data()
      • create_process_specific_iterator()
      • create_data_iterator()
    • maxtext.input_pipeline.input_pipeline_utils module
      • normalize_features()
      • get_tokenizer()
      • truncate_to_max_allowable_length()
      • shift_data_by_truncation()
      • add_segmentation_and_position()
      • TokenizeOp()
      • reformat_prompt()
      • reformat_response()
      • merge_image_columns()
      • pre_process_image_sft()
      • prepare_text_for_image_fusion()
      • combine_columns()
      • is_conversational()
      • extract_token_ids()
      • verify_chat_template_generation_prompt_logic()
      • apply_chat_template()
      • tokenization()
      • SFTPromptMasking
      • SFTPromptMaskingVision
      • HFNormalizeFeatures
      • HFDataSource
      • GCSTFRecordIterDataset
      • make_tfrecord_iter_dataset()
      • ParseFeatures
      • NormalizeFeatures
      • KeepFeatures
      • Rekey
      • ReformatPacking
      • PadOrTrimToMaxLength
      • ExtractImagesAndMasks
      • FoldImagesIntoBatch
      • shift_right()
      • shift_left()
      • shift_and_refine()
      • ShiftData
      • ComputeQwen3OmniPositions
    • maxtext.input_pipeline.instruction_data_processing module
      • load_data_template_from_file()
      • load_chat_template_from_file()
      • get_template_placeholders()
      • extract_reasoning_and_answer()
      • math_qa_formatting()
      • load_formatter()
      • convert_to_conversational_format()
    • maxtext.input_pipeline.multihost_dataloading module
      • MultiHostDataLoadIterator
      • RemoteIteratorWrapper
    • maxtext.input_pipeline.olmo_data module
      • OlmoNpyFileEntry
      • OlmoNpyIndex
      • load_index()
      • global_to_local()
      • compute_fingerprint()
      • parse_npy_header()
      • read_npy_header_from_path()
      • read_raw_metadata_from_path()
      • has_npy_magic()
      • build_index()
      • RepetitionTuple
      • find_periodic_sequences()
      • is_clean_instance()
    • maxtext.input_pipeline.olmo_data_grain module
      • OlmoNpyDataSource
      • OlmoIndexSampler
      • NgramFilterTransform
      • ShiftToInputsTargets
      • make_olmo_grain_data_loader()
    • maxtext.input_pipeline.olmo_grain_data_processing module
      • make_olmo_grain_train_iterator()
      • make_olmo_grain_eval_iterator()
    • maxtext.input_pipeline.synthetic_data_processing module
      • SyntheticDataIterator
      • PlaceHolderDataIterator
    • maxtext.input_pipeline.tfds_data_processing module
      • get_datasets()
      • preprocessing_pipeline()
      • make_tfds_train_iterator()
      • make_tfds_eval_iterator()
    • maxtext.input_pipeline.tfds_data_processing_c4_mlperf module
      • rekey()
      • reduce_concat_tokens()
      • split_tokens()
      • split_tokens_to_targets_length()
      • get_dataset()
      • format_fn()
      • preprocess_train_dataset()
      • preprocess_eval_dataset()
      • make_c4_mlperf_train_iterator()
      • make_c4_mlperf_eval_iterator()
    • maxtext.input_pipeline.tokenizer module
      • TikTokenTokenizer
      • SentencePieceTokenizer
      • HFTokenizer
      • build_tokenizer()
• maxtext.integration package
  • Subpackages
    • maxtext.integration.tunix package
      • Subpackages
      • Submodules
    • maxtext.integration.vllm package
      • Subpackages
      • Submodules
• maxtext.kernels package
  • Subpackages
    • maxtext.kernels.attention package
      • Submodules
    • maxtext.kernels.megablox package
      • Submodules
  • Submodules
    • maxtext.kernels.gather_reduce_pallas module
      • sc_gather_reduce()
    • maxtext.kernels.gather_reduce_sc module
      • VectorTypeHelper
      • sc_gather_reduce()
    • maxtext.kernels.sort_activations module
• maxtext.layers package
  • Submodules
    • maxtext.layers.attention_mla module
      • Indexer
      • mla_as_linen()
      • MLA
    • maxtext.layers.attention_op module
      • validate_compute_axis_order()
      • apply_mask_to_logits()
      • validate_gpu_flash_attention()
      • ChunkedCausalMask
      • attention_op_as_linen()
      • AttentionOp
      • LoadBalancedCausalMask
    • maxtext.layers.attentions module
      • L2Norm
      • l2_norm_as_linen()
      • attention_as_linen()
      • Attention
    • maxtext.layers.decoders module
      • DecoderLayer
      • SequentialBlockDecoderLayers
      • deepstack_process()
      • Decoder
    • maxtext.layers.embeddings module
      • embed_as_linen()
      • Embed
      • attend_on_embedding()
      • rotary_embedding_as_linen()
      • RotaryEmbedding
      • llama_rotary_embedding_as_linen()
      • partial_rotary_embedding_as_linen()
      • PartialRotaryEmbedding
      • Gemma4PartialRotaryEmbedding
      • LLaMARotaryEmbedding
      • yarn_rotary_embedding_as_linen()
      • YarnRotaryEmbedding
      • positional_embedding_as_linen()
      • PositionalEmbedding
      • llama_vision_rotary_embedding_as_linen()
      • LlamaVisionRotaryEmbedding
      • Qwen3OmniMoeVisionRotaryEmbedding
      • qwen3omnimoe_vision_pos_embed_interpolate_as_linen()
      • Qwen3OmniMoeVisionPosEmbedInterpolate
      • 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
      • NgramHashMapping
      • StaticWrapper
      • MultiHeadEmbedding
      • ShortConv
      • Engram
    • maxtext.layers.initializers module
      • nd_dense_init()
      • variable_to_logically_partitioned()
    • maxtext.layers.learn_to_init_layer module
      • LearnToInitDecoderLayer
      • LearnToInitDense
      • calculate_attn_weight()
      • apply_lti_model_update()
    • maxtext.layers.linears module
      • normalize_axes()
      • canonicalize_tuple()
      • DenseGeneral
      • dense_general()
      • Dropout
      • MlpBlock
      • mlp_block()
    • maxtext.layers.mhc module
      • get_functions()
      • sinkhorn()
      • ManifoldConstrainedHyperConnections
    • maxtext.layers.moe module
      • get_batchsplit_init_kernel_axes()
      • random_routing()
      • calculate_load_balance_updates()
      • GateLogit
      • RoutedMoE
      • RoutedAndSharedMoE
      • 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
      • MultiTokenPredictionBlock
      • calculate_mtp_loss()
      • calculate_mtp_acceptance_rate()
      • multi_token_prediction_block_as_linen()
    • maxtext.layers.nnx_decoders module
      • NNXDecoderLayer
      • deepstack_process()
      • NNXDecoder
      • 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
      • linen_rngs_dict()
      • ToLinen
      • to_linen()
      • to_linen_class()
    • maxtext.layers.normalizations module
      • RMSNorm
      • GlobalRMSNorm
      • Qwen3NextRMSNorm()
      • Qwen3NextRMSNormGated
      • rms_norm()
      • l2norm()
    • maxtext.layers.pipeline module
      • PipelineBase
      • Pipeline
      • CircularPipeline
      • create_pipeline()
    • maxtext.layers.pipeline_deprecated module
      • Pipeline
    • maxtext.layers.quantizations module
      • Quantization
      • AqtQuantization
      • QwixQuantization
      • QwixDotGeneral
      • QwixEinsum
      • Fp8Quantization
      • Fp8Einsum
      • NANOOFp8Quantization
      • ConstantBoundConfig
      • PerTensorScales
      • in_convert_mode()
      • in_serve_mode()
      • get_quant_mode()
      • configure_quantization()
      • match_aqt_and_unquantized_param()
      • remove_quantized_params()
      • configure_kv_quant()
      • NvidaFp8Provider
      • NANOOFp8Provider
      • get_fp8_full_qwix_rule_w_sparsity()
      • get_quantization_rule()
      • get_qt_provider()
      • maybe_quantize_model()
      • manual_quantize()
      • TransformerEngineQuantization
    • maxtext.layers.train_state_nnx module
      • TrainStateNNX
• maxtext.models package
  • Submodules
    • maxtext.models.deepseek module
      • DeepSeekGenericLayer
      • DeepSeekDenseLayer
      • DeepSeekMoELayer
    • 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
    • maxtext.models.gemma2 module
      • Gemma2DecoderLayer
    • maxtext.models.gemma3 module
      • get_attention_type()
      • get_query_pre_attn_scalar()
      • Gemma3DecoderLayer
      • Gemma3ScannableBlock
      • MlpBlockViT
      • Encoder1DBlock
      • Encoder
      • Einsum
      • VisionEmbedder
      • visionembedder_as_linen()
      • VisionExit
      • 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
      • Gemma4EncoderBlock
      • Gemma4VisionEncoderLayer
      • Gemma4VisionProjector
      • gemma4_vision_encoder_as_linen()
    • maxtext.models.gpt3 module
      • Gpt3LayerNorm
      • gpt3_layer_norm()
      • Gpt3MultiHeadAttention
      • Gpt3DecoderLayer
    • maxtext.models.gpt_oss module
      • get_attention_type()
      • GptOssDecoderLayer
      • GptOssScannableBlock
    • maxtext.models.llama2 module
      • LlamaDecoderLayer
      • LlamaLTIDecoderLayer
    • maxtext.models.llama4 module
      • Llama4UnfoldConvolution
      • pixel_shuffle()
      • Llama4VisionMLP
      • Llama4VisionMLP2
      • Llama4VisionPixelShuffleMLP
      • Llama4MultiModalProjector
      • llama4multimodalprojector_as_linen()
      • determine_is_nope_layer()
      • determine_is_moe_layer()
      • Llama4DecoderLayer
      • Llama4ScannableBlock
      • Llama4VisionEncoderLayer
      • Llama4VisionEncoder
      • Llama4VisionModel
      • llama4visionmodel_as_linen()
    • maxtext.models.mistral module
      • MistralDecoderLayer
    • maxtext.models.mixtral module
      • MixtralDecoderLayer
    • maxtext.models.models module
      • TransformerLinenPure
      • transformer_as_linen()
      • TransformerLinen
      • Transformer
    • maxtext.models.olmo3 module
      • get_attention_type()
      • Olmo3DecoderLayer
      • Olmo3ScannableBlock
    • maxtext.models.qwen2 module
      • AttentionWithNorm
      • Qwen2DecoderLayer
    • maxtext.models.qwen3 module
      • naive_jax_chunk_gated_delta_rule()
      • jax_chunk_gated_delta_rule()
      • Qwen3NextGatedDeltaNet
      • Qwen3NextFullAttention
      • Qwen3NextSparseMoeBlock
      • Qwen3NextScannableBlock
      • Qwen3NextDecoderLayer
      • AttentionWithNorm
      • Qwen3DecoderLayer
      • Qwen3MoeDecoderLayer
      • Qwen3OmniMoeVisionPatchMerger
      • Qwen3OmniMoeVisionMLP
      • Qwen3OmniMoeVisionPatchEmbed
      • Qwen3OmniMoeVisionAttention
      • Qwen3OmniMoeVisionBlock
      • Qwen3OmniMoeVisionEncoder
      • Qwen3OmniMoeVisionProjector
      • qwen3omni_visionencoder_as_linen()
      • qwen3omni_visionprojector_as_linen()
      • Qwen3OmniAudioEncoderLayer
      • Qwen3OmniAudioEncoder
      • Qwen3OmniAudioProjector
      • qwen3omni_audioencoder_as_linen()
      • qwen3omni_audioprojector_as_linen()
    • maxtext.models.qwen3_5 module
      • Qwen3_5GatedDeltaNet
      • Qwen3_5FullAttention
      • Qwen3_5SparseMoEBlock
      • Qwen3_5ScannableBlock
      • Qwen3_5DecoderLayer
    • maxtext.models.qwen3_custom module
      • Qwen3CustomAttention
      • Qwen3CustomMoeDecoderLayer
    • maxtext.models.simple_layer module
      • SimpleDecoderLayer
      • SimpleMlpDecoderLayer
• maxtext.multimodal package
  • Submodules
    • maxtext.multimodal.processor module
      • preprocess_mm_data()
      • preprocess_image_for_training()
      • get_image_offsets()
      • reformat_prompt()
      • reformat_response()
      • prepare_text_for_image_fusion()
      • get_dummy_image_shape_for_init()
      • get_dummy_audio_shape_for_init()
      • get_bidirectional_mask_vision()
      • get_bidirectional_mask_audio()
    • maxtext.multimodal.processor_gemma3 module
      • Gemma3PreprocessorOutput
      • preprocess_mm_data_gemma3()
      • get_image_offsets_gemma3()
      • reformat_prompt_gemma3()
      • insert_sequence()
      • add_extra_tokens_for_images_gemma3()
      • get_dummy_image_shape_for_init_gemma3()
    • maxtext.multimodal.processor_gemma4 module
      • Gemma4PreprocessorOutput
      • preprocess_mm_data_gemma4()
      • get_image_offsets_gemma4()
      • reformat_prompt_gemma4()
      • insert_sequence()
      • add_extra_tokens_for_images_gemma4()
      • get_dummy_image_shape_for_init_gemma4()
    • maxtext.multimodal.processor_llama4 module
      • Llama4PreprocessorOutput
      • get_factors()
      • find_supported_resolutions()
      • get_best_resolution()
      • pad_to_best_fit_jax()
      • pad_to_max_tiles()
      • split_to_tiles()
      • preprocess_mm_data_llama4()
      • get_num_tokens_for_this_image()
      • get_image_offsets_llama4()
      • reformat_prompt_llama4()
      • get_tokens_for_this_image()
      • add_extra_tokens_for_images_llama4()
      • get_dummy_image_shape_for_init_llama4()
    • maxtext.multimodal.processor_qwen3_omni module
      • Qwen3OmniPreprocessorOutput
      • smart_resize()
      • pre_process_qwen3_image()
      • calculate_video_frame_range()
      • smart_nframes()
      • preprocess_video()
      • pre_process_audio_qwen3_omni()
      • preprocess_mm_data_qwen3_omni()
      • add_extra_tokens_for_qwen3_omni()
      • get_dummy_image_shape_for_init_qwen3_omni()
      • get_dummy_audio_shape_for_init_qwen3_omni()
      • get_llm_pos_ids_for_vision()
      • get_chunked_index()
      • get_rope_index()
      • reformat_prompt_qwen3_omni()
      • get_mm_offsets_qwen3_omni()
    • maxtext.multimodal.utils module
      • PreprocessorOutput
      • convert_to_RGB()
      • load_image_from_path()
      • normalize_images()
      • merge_mm_embeddings()
      • amplitude_to_db()
      • power_to_db()
      • spectrogram()
      • hertz_to_mel()
      • mel_to_hertz()
      • mel_filter_bank()
      • window_function()
      • load_audio()