class_docstr: Wrapper class for Mesh-TF models.
class_docstr_tok: ["Wrapper","class","for","Mesh-TF","models","."]
docstr: Constructor for MtfModel class.
  
  Args:
    model_dir: str, directory to save the model.
    tpu: str, the TPU address to use.
    tpu_job_name: str, name of the TPU worker binary.
    tpu_zone: str, GCE zone where the Cloud TPU is located
    gcp_project: str, project name for the Cloud TPU-enabled project.
    tpu_topology: str, e.g. "2x2" or "v2-8".
    model_parallelism: integer, the number of cores per model replica.
    batch_size: An integer or a (method, value) pair to pass to
      compute_batch_size(). Note that this is the global batch size and not
      the per-shard batch size.
    sequence_length: an integer or a dict from feature-key to integer
      the (packed) sequence length, e.g. {"inputs": 512, "targets": 128}
    model_type: str, a model type from mesh tf models.
    layout_rules: an input to mtf.convert_to_layout_rules()
    mesh_shape: an mtf.Shape or string (e.g., "model:2,batch:4") specifying
      how the data/model should be parallelized. If None (default), the mesh
      shape will be constructed using the supplied `tpu_topology` and
      `model_parallelism` arguments.
    mesh_devices: a list of strings, the device names to use for each mesh
      slice. Only required for GPU.
    autostack: boolean, internally combine variables.
    learning_rate_schedule: an optional function taking the scalar name
      argument `step` and the numeric argument `total_train_steps` and return
      the scalar learning rate.
    keep_checkpoint_max: an integer, maximum number of checkpoints to keep.
    save_checkpoints_steps: an integer, steps per checkpoint.
    optimizer: a class extending optimize.Optimizer, required for training.
    predict_fn: an optional function that can be used to override the default
      transformer prediction behavior. Must return a tensor of shape
      [batch_dim, length_dim] that will be the prediction for each example.
      Must accept the following arguments:
        - model: a Unitransformer or Bitransformer
        - features: a dict representing an example. Every value will be an
          mtf.Tensor with shape [batch_dim, length_dim].
        - variable_dtype: an mtf.VariableDType
    variable_filter: a str, a variable will only be trained if its name
      matches this regex. If None (default), train all trainable variables.
    ensemble_inputs: an integer, see `train_model` docstring for details.
    iterations_per_loop: integer, steps per train loop
    extra_gin_bindings: an optional list of strings, extra gin bindings to
      pass to `gin.parse_config` after loading the operative config.
docstr_tok: ["Constructor","for","MtfModel","class",".","Args",":","model_dir",":","str",",","directory","to","save","the","model",".","tpu",":","str",",","the","TPU","address","to","use",".","tpu_job_name",":","str",",","name","of","the","TPU","worker","binary",".","tpu_zone",":","str",",","GCE","zone","where","the","Cloud","TPU","is","located","gcp_project",":","str",",","project","name","for","the","Cloud","TPU-enabled","project",".","tpu_topology",":","str",",","e.g",".","``","2x2","''","or","``","v2-8","''",".","model_parallelism",":","integer",",","the","number","of","cores","per","model","replica",".","batch_size",":","An","integer","or","a","(","method",",","value",")","pair","to","pass","to","compute_batch_size","(",")",".","Note","that","this","is","the","global","batch","size","and","not","the","per-shard","batch","size",".","sequence_length",":","an","integer","or","a","dict","from","feature-key","to","integer","the","(","packed",")","sequence","length",",","e.g",".","{","``","inputs","''",":","512",",","``","targets","''",":","128","}","model_type",":","str",",","a","model","type","from","mesh","tf","models",".","layout_rules",":","an","input","to","mtf.convert_to_layout_rules","(",")","mesh_shape",":","an","mtf.Shape","or","string","(","e.g.",",","``","model:2",",","batch:4","''",")","specifying","how","the","data\/model","should","be","parallelized",".","If","None","(","default",")",",","the","mesh","shape","will","be","constructed","using","the","supplied","`","tpu_topology","`","and","`","model_parallelism","`","arguments",".","mesh_devices",":","a","list","of","strings",",","the","device","names","to","use","for","each","mesh","slice",".","Only","required","for","GPU",".","autostack",":","boolean",",","internally","combine","variables",".","learning_rate_schedule",":","an","optional","function","taking","the","scalar","name","argument","`","step","`","and","the","numeric","argument","`","total_train_steps","`","and","return","the","scalar","learning","rate",".","keep_checkpoint_max",":","an","integer",",","maximum","number","of","checkpoints","to","keep",".","save_checkpoints_steps",":","an","integer",",","steps","per","checkpoint",".","optimizer",":","a","class","extending","optimize.Optimizer",",","required","for","training",".","predict_fn",":","an","optional","function","that","can","be","used","to","override","the","default","transformer","prediction","behavior",".","Must","return","a","tensor","of","shape","[","batch_dim",",","length_dim","]","that","will","be","the","prediction","for","each","example",".","Must","accept","the","following","arguments",":","-","model",":","a","Unitransformer","or","Bitransformer","-","features",":","a","dict","representing","an","example",".","Every","value","will","be","an","mtf.Tensor","with","shape","[","batch_dim",",","length_dim","]",".","-","variable_dtype",":","an","mtf.VariableDType","variable_filter",":","a","str",",","a","variable","will","only","be","trained","if","its","name","matches","this","regex",".","If","None","(","default",")",",","train","all","trainable","variables",".","ensemble_inputs",":","an","integer",",","see","`","train_model","`","docstring","for","details",".","iterations_per_loop",":","integer",",","steps","per","train","loop","extra_gin_bindings",":","an","optional","list","of","strings",",","extra","gin","bindings","to","pass","to","`","gin.parse_config","`","after","loading","the","operative","config","."]
code: def __init__(
      self,
      model_dir,
      tpu,
      tpu_job_name=None,
      tpu_zone=None,
      gcp_project=None,
      tpu_topology="v2-8",
      model_parallelism=8,
      batch_size=("sequences_per_batch", 1),
      sequence_length=None,
      model_type="bitransformer",
      layout_rules="ensemble:ensemble,batch:batch,d_ff:model,heads:model,vocab:model,experts:batch",
      mesh_shape=None,
      mesh_devices=None,
      autostack=True,
      learning_rate_schedule=None,
      keep_checkpoint_max=None,
      save_checkpoints_steps=5000,
      optimizer=None,
      predict_fn=None,
      variable_filter=None,
      ensemble_inputs=None,
      iterations_per_loop=100,
      extra_gin_bindings=None,
  ):
      """Constructor for MtfModel class.
  
      Args:
        model_dir: str, directory to save the model.
        tpu: str, the TPU address to use.
        tpu_job_name: str, name of the TPU worker binary.
        tpu_zone: str, GCE zone where the Cloud TPU is located
        gcp_project: str, project name for the Cloud TPU-enabled project.
        tpu_topology: str, e.g. "2x2" or "v2-8".
        model_parallelism: integer, the number of cores per model replica.
        batch_size: An integer or a (method, value) pair to pass to
          compute_batch_size(). Note that this is the global batch size and not
          the per-shard batch size.
        sequence_length: an integer or a dict from feature-key to integer
          the (packed) sequence length, e.g. {"inputs": 512, "targets": 128}
        model_type: str, a model type from mesh tf models.
        layout_rules: an input to mtf.convert_to_layout_rules()
        mesh_shape: an mtf.Shape or string (e.g., "model:2,batch:4") specifying
          how the data/model should be parallelized. If None (default), the mesh
          shape will be constructed using the supplied `tpu_topology` and
          `model_parallelism` arguments.
        mesh_devices: a list of strings, the device names to use for each mesh
          slice. Only required for GPU.
        autostack: boolean, internally combine variables.
        learning_rate_schedule: an optional function taking the scalar name
          argument `step` and the numeric argument `total_train_steps` and return
          the scalar learning rate.
        keep_checkpoint_max: an integer, maximum number of checkpoints to keep.
        save_checkpoints_steps: an integer, steps per checkpoint.
        optimizer: a class extending optimize.Optimizer, required for training.
        predict_fn: an optional function that can be used to override the default
          transformer prediction behavior. Must return a tensor of shape
          [batch_dim, length_dim] that will be the prediction for each example.
          Must accept the following arguments:
            - model: a Unitransformer or Bitransformer
            - features: a dict representing an example. Every value will be an
              mtf.Tensor with shape [batch_dim, length_dim].
            - variable_dtype: an mtf.VariableDType
        variable_filter: a str, a variable will only be trained if its name
          matches this regex. If None (default), train all trainable variables.
        ensemble_inputs: an integer, see `train_model` docstring for details.
        iterations_per_loop: integer, steps per train loop
        extra_gin_bindings: an optional list of strings, extra gin bindings to
          pass to `gin.parse_config` after loading the operative config.
      """
      mesh_shape = mesh_shape or (
          mtf_utils.tpu_mesh_shape(tpu_topology, model_parallelism)
          if tpu
          else ""
      )
  
      sequence_length = sequence_length or {
          "inputs": 512,
          "targets": 512,
      }
  
      if isinstance(sequence_length, int):
          sequence_length = {
              "inputs": sequence_length,
              "targets": sequence_length,
          }
      self._learning_rate_schedule = (
          learning_rate_schedule
          or learning_rate_schedules.learning_rate_schedule_noam
      )
  
      self._optimizer = optimizer or optimize.AdafactorOptimizer
  
      self._sequence_length = sequence_length
      self._model_dir = model_dir
      self._model_type = model_type
      self._ensemble_inputs = ensemble_inputs
  
      self._layout_rules = mtf.convert_to_layout_rules(layout_rules)
      self._mesh_shape = mtf.convert_to_shape(mesh_shape)
      self._mesh_devices = mesh_devices
  
      self._autostack = autostack
      self._keep_checkpoint_max = keep_checkpoint_max
      self._save_checkpoints_steps = save_checkpoints_steps
      self._predict_fn = predict_fn
      self._variable_filter = variable_filter
      self._ensemble_inputs = ensemble_inputs
      self._iterations_per_loop = iterations_per_loop
  
      self._cluster = (
          tf.distribute.cluster_resolver.TPUClusterResolver(
              tpu, zone=tpu_zone, project=gcp_project
          )
          if tpu
          else None
      )
      self._tpu = tpu
      self._tpu_job_name = tpu_job_name
      self._estimator = None
  
      # Must be called after _sequence_length, _mesh_shape, and _layout_rules are
      # set.
      self.batch_size = batch_size
  
      self._gin_bindings = extra_gin_bindings
  
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 1 to create a new MtfModel object from inherited base classes: T5Model with example: obj = MtfModel(model_dir, tpu, tpu_job_name, tpu_zone, gcp_project, tpu_topology, model_parallelism, batch_size, sequence_length, model_type, layout_rules, mesh_shape, mesh_devices, autostack, learning_rate_schedule, keep_checkpoint_max, save_checkpoints_steps, optimizer, predict_fn, variable_filter, ensemble_inputs, iterations_per_loop, extra_gin_bindings)