maxtext.integration.tunix.tunix_adapter module

maxtext.integration.tunix.tunix_adapter module#

Adapter for integrating MaxText Transformer models with Tunix.

This module provides the TunixMaxTextAdapter class, which wraps a MaxText Transformer model to expose a call signature compatible with Tunix Trainers. It also handles weight mapping for compatibility with Hugging Face models.

class maxtext.integration.tunix.tunix_adapter.TunixMaxTextAdapter(*args, **kwargs)[source]#

Bases: Module

Adapter exposing Tunix Trainer call signature over a Transformer model.

Parameters:

  • args (Any)

  • kwargs (Any)

Return type:

Any

to_hf_mappings()[source]#
to_hf_transpose_keys()[source]#
to_hf_hook_fns()[source]#
lora_to_hf_mappings()[source]#