class_docstr: Toy model to test mesh_tensorflow.
class_docstr_tok: ["Toy","model","to","test","mesh_tensorflow","."]
code: def estimator_spec_predict(self, features, mesh, mesh_impl, use_tpu):
      mtf_samples = mtf.anonymize(self.sample(features, mesh))
      lowering = mtf.Lowering(mesh.graph, {mesh: mesh_impl})
      outputs = lowering.export_to_tf_tensor(mtf_samples)
      if self.has_input:
          ndims = len(outputs.shape.as_list())
          actual_batch_size = tf.shape(features["inputs"])[0]
          outputs = tf.slice(
              outputs,
              [0] * ndims,
              [actual_batch_size] + [-1] * (ndims - 1),
          )
      predictions = {"outputs": outputs}
      if features.get("infer_targets") is not None:
          predictions["infer_targets"] = features["infer_targets"]
  
      if features.get("inputs") is not None:
          predictions["inputs"] = features["inputs"]
  
      if use_tpu:
          t2t_model.remove_summaries()
          return tpu_estimator.TPUEstimatorSpec(
              mode=tf_estimator.ModeKeys.PREDICT,
              predictions=predictions,
              prediction_hooks=[mtf.MtfRestoreHook(lowering)],
          )
      else:
          return tf_estimator.EstimatorSpec(
              tf_estimator.ModeKeys.PREDICT,
              predictions=predictions,
              prediction_hooks=[mtf.MtfRestoreHook(lowering)],
          )
  
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 3 for calling the MtfModel obj's underlying member method code with example usage: obj.estimator_spec_predict(features, mesh, mesh_impl, use_tpu) and returns: tpu_estimator, tf_estimator