class_docstr: Toy model to test mesh_tensorflow.
class_docstr_tok: ["Toy","model","to","test","mesh_tensorflow","."]
code: def estimator_model_fn(
      cls,
      hparams,
      features,
      labels,
      mode,
      config=None,
      params=None,
      decode_hparams=None,
      use_tpu=False,
  ):
      hparams = hparams_lib.copy_hparams(hparams)
      hparams.use_tpu = use_tpu
      # merge decode_hparams into hparams if present
      if (
          mode == tf_estimator.ModeKeys.PREDICT
          and decode_hparams is not None
      ):
          for k, v in six.iteritems(decode_hparams.values()):
              if hasattr(hparams, k) and getattr(hparams, k) != v:
                  tf.logging.warning(
                      "Overriding hparams.%s with %s from decode_hparams"
                      % (k, v)
                  )
              setattr(hparams, k, v)
  
      # Instantiate model
      data_parallelism = None
      if not use_tpu and config:
          data_parallelism = config.data_parallelism
      model = cls(
          hparams,
          mode,
          data_parallelism=data_parallelism,
          decode_hparams=decode_hparams,
      )
  
      global_step = tf.train.get_global_step()
  
      mesh_shape = mtf.convert_to_shape(hparams.mesh_shape)
      layout_rules = mtf.convert_to_layout_rules(hparams.layout)
      if use_tpu:
          ctx = params["context"]
          num_hosts = ctx.num_hosts
          host_placement_fn = ctx.tpu_host_placement_function
          device_list = [
              host_placement_fn(host_id=t) for t in range(num_hosts)
          ]
          # TODO(ylc): Better estimation of replica cache size?
          replica_cache_size = 300 * 1000000  # 300M per replica
          # Worker 0 caches all the TPU binaries.
          worker0_mem = replica_cache_size * ctx.num_replicas
          devices_memeory_usage = [worker0_mem] + [0] * (num_hosts - 1)
          var_placer = mtf.utils.BalancedVariablePlacer(
              device_list, devices_memeory_usage
          )
          mesh_devices = [""] * mesh_shape.size
          mesh_impl = mtf.simd_mesh_impl.SimdMeshImpl(
              mesh_shape,
              layout_rules,
              mesh_devices,
              ctx.device_assignment,
          )
      else:
          var_placer = None
          if (
              data_parallelism is None
              or len(data_parallelism.ps_devices) == 1
          ):
              mesh_devices = [""] * mesh_shape.size
          else:
              assert len(data_parallelism.ps_devices) == mesh_shape.size
              mesh_devices = data_parallelism.ps_devices
          mesh_impl = mtf.placement_mesh_impl.PlacementMeshImpl(
              mesh_shape, layout_rules, mesh_devices
          )
  
      graph = mtf.Graph()
      mesh = mtf.Mesh(graph, "my_mesh", var_placer)
      # PREDICT mode
      if mode == tf_estimator.ModeKeys.PREDICT:
          return model.estimator_spec_predict(
              features, mesh, mesh_impl, use_tpu
          )
  
      logits, loss = model.mtf_model_fn(features, mesh)
      if use_tpu and logits is not None:
          logits = mtf.anonymize(logits)
  
      # TRAIN mode
      if mode == tf_estimator.ModeKeys.TRAIN:
          var_grads = mtf.gradients(
              [loss], [v.outputs[0] for v in graph.trainable_variables]
          )
          lr = learning_rate.learning_rate_schedule(hparams)
          tf.summary.scalar("learning_rate", lr)
          mtf_lr = mtf.import_tf_tensor(
              mesh,
              tf.convert_to_tensor(lr, dtype=tf.float32),
              mtf.Shape([]),
          )
          optimizer = mtf.optimize.make_optimizer(hparams, mtf_lr)
          update_ops = optimizer.apply_grads(
              var_grads, graph.trainable_variables
          )
  
      lowering = mtf.Lowering(graph, {mesh: mesh_impl})
  
      tf_loss = lowering.export_to_tf_tensor(loss)
      tf_loss = tf.to_float(tf_loss)
      if logits and mode != tf_estimator.ModeKeys.TRAIN:
          tf_logits = lowering.export_to_tf_tensor(logits)
  
      if mode == tf_estimator.ModeKeys.TRAIN:
          tf_update_ops = [
              lowering.lowered_operation(op) for op in update_ops
          ]
          tf_update_ops.append(tf.assign_add(global_step, 1))
          # tf.logging.info("tf_update_ops: {}".format(tf_update_ops))
          train_op = tf.group(tf_update_ops)
  
      with mtf.utils.outside_all_rewrites():
          # Copy master variables to slices. Must be called first.
          restore_hook = mtf.MtfRestoreHook(lowering)
          saver = tf.train.Saver(
              tf.global_variables(),
              sharded=True,
              max_to_keep=10,
              keep_checkpoint_every_n_hours=2,
              defer_build=False,
              save_relative_paths=True,
          )
          tf.add_to_collection(tf.GraphKeys.SAVERS, saver)
          saver_listener = mtf.MtfCheckpointSaverListener(lowering)
          saver_hook = tf.train.CheckpointSaverHook(
              hparams.model_dir,
              save_steps=1000,
              saver=saver,
              listeners=[saver_listener],
          )
  
      # EVAL mode
      if mode == tf_estimator.ModeKeys.EVAL:
          tf_logits = lowering.export_to_tf_tensor(logits)
          return model.estimator_spec_eval(
              features,
              tf_logits,
              labels,
              tf_loss,
              restore_hook,
              use_tpu,
          )
  
      if use_tpu:
          # TPU host call. Important: need to be called before remove_summaries()
          if hparams.tpu_enable_host_call:
              host_call = t2t_model.create_host_call(hparams.model_dir)
          else:
              host_call = None
  
          if hparams.warm_start_from:
  
              def scaffold_fn():
                  t2t_model.initialize_from_ckpt(
                      ckpt_dir=hparams.warm_start_from, hparams=hparams
                  )
                  return tf.train.Scaffold()
  
          else:
              scaffold_fn = None
  
          t2t_model.remove_summaries()
          return tpu_estimator.TPUEstimatorSpec(
              mode=tf_estimator.ModeKeys.TRAIN,
              loss=tf_loss,
              train_op=train_op,
              host_call=host_call,
              training_hooks=[restore_hook, saver_hook],
              scaffold_fn=scaffold_fn,
          )
      else:
          if hparams.warm_start_from:
              t2t_model.initialize_from_ckpt(
                  ckpt_dir=hparams.warm_start_from, hparams=hparams
              )
          return tf_estimator.EstimatorSpec(
              tf_estimator.ModeKeys.TRAIN,
              loss=tf_loss,
              train_op=train_op,
              training_chief_hooks=[restore_hook, saver_hook],
          )
  
file_path: tensor2tensor/tensor2tensor/utils/mtf_model.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import mesh_tensorflow
  import six
  from tensor2tensor.utils import hparams_lib
  from tensor2tensor.utils import learning_rate
  from tensor2tensor.utils import metrics
  from tensor2tensor.utils import t2t_model
  import tensorflow.compat.v1
  from tensorflow.compat.v1 import estimator
  from tensorflow.contrib.tpu.python.tpu import tpu_estimator
label_desc: Use image node_id 1 for calling the MtfModel obj's underlying member method code with example usage: obj.estimator_model_fn(cls, hparams, features, labels, mode, config, params, decode_hparams, use_tpu) and returns: model, model, tpu_estimator, tf_estimator, tf