class_docstr: Wrapper class for Mesh-TF models.
class_docstr_tok: ["Wrapper","class","for","Mesh-TF","models","."]
docstr: Train the model on the given Mixture or Task.
  
  Args:
    mixture_or_task_name: str, the name of the Mixture or Task to train on.
      Must be pre-registered in the global `TaskRegistry` or
      `MixtureRegistry.`
    steps: int, the total number of steps to train for.
    init_checkpoint: a string, if not None then read in variables from this
      checkpoint path when initializing variables. Will only initialize
      variables that appear both in the current graph and the checkpoint.
    split: str, the mixture/task split to train on.
docstr_tok: ["Train","the","model","on","the","given","Mixture","or","Task",".","Args",":","mixture_or_task_name",":","str",",","the","name","of","the","Mixture","or","Task","to","train","on",".","Must","be","pre-registered","in","the","global","`","TaskRegistry","`","or","`","MixtureRegistry.","`","steps",":","int",",","the","total","number","of","steps","to","train","for",".","init_checkpoint",":","a","string",",","if","not","None","then","read","in","variables","from","this","checkpoint","path","when","initializing","variables",".","Will","only","initialize","variables","that","appear","both","in","the","current","graph","and","the","checkpoint",".","split",":","str",",","the","mixture\/task","split","to","train","on","."]
code: def train(
      self,
      mixture_or_task_name,
      steps,
      init_checkpoint=None,
      split="train",
  ):
      """Train the model on the given Mixture or Task.
  
      Args:
        mixture_or_task_name: str, the name of the Mixture or Task to train on.
          Must be pre-registered in the global `TaskRegistry` or
          `MixtureRegistry.`
        steps: int, the total number of steps to train for.
        init_checkpoint: a string, if not None then read in variables from this
          checkpoint path when initializing variables. Will only initialize
          variables that appear both in the current graph and the checkpoint.
        split: str, the mixture/task split to train on.
      """
      vocabulary = mesh_transformer.get_vocabulary(mixture_or_task_name)
      dataset_fn = functools.partial(
          mesh_transformer.mesh_train_dataset_fn,
          mixture_or_task_name=mixture_or_task_name,
      )
      mtf_utils.train_model(
          self.estimator(vocabulary, init_checkpoint),
          vocabulary,
          self._sequence_length,
          self.batch_size,
          dataset_fn,
          steps,
          self._ensemble_inputs,
          dataset_split=split,
      )
  
imports: import functools
  import os
  import gin
  import gin.tf
  import mesh_tensorflow
  from mesh_tensorflow import optimize
  from mesh_tensorflow.transformer import dataset
  from mesh_tensorflow.transformer import learning_rate_schedules
  from mesh_tensorflow.transformer import utils
  from t5.models import mesh_transformer
  from t5.models import utils
  from t5.models.t5_model import T5Model
  import tensorflow.compat.v1
label_desc: Use image node_id 5 for calling the MtfModel obj's underlying member method code with example usage: obj.train(mixture_or_task_name, steps, init_checkpoint, split) without return types