<commit_before># coding=utf-8
# Copyright 2019 The Mesh TensorFlow Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

r"""Utilities for running training and inference."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import functools
import math
import os
import re

import gin
import gin.tf

import mesh_tensorflow as mtf
from mesh_tensorflow.transformer import dataset as transformer_dataset
from mesh_tensorflow.transformer import transformer
import numpy as np
import six
import tensorflow as tf
import tensorflow_datasets as tfds

from tensorflow.contrib.tpu.python.tpu import tpu_config
from tensorflow.contrib.tpu.python.tpu import tpu_estimator

tf.flags.DEFINE_multi_string("gin_file", None,
                             "List of paths to the config files.")
tf.flags.DEFINE_multi_string(
    "gin_param", None, "Newline separated list of Gin parameter bindings.")
FLAGS = tf.flags.FLAGS

_DEFAULT_CONFIG_FILE = "./gin/defaults.gin"

# List of features used by model.
_INPUT_FEATURES = [
    "inputs", "inputs_position", "inputs_segmentation", "targets",
    "targets_position", "targets_segmentation", "targets_subsegmentation"
]


def parse_gin_defaults_and_flags():
  """Parses all default gin files and those provides via flags."""
  # Set up the default values for the configurable parameters. These values will
  # be overridden by any user provided gin files/parameters.
  gin.parse_config_file(
      os.path.join(os.path.dirname(__file__), _DEFAULT_CONFIG_FILE))
  gin.parse_config_files_and_bindings(FLAGS.gin_file, FLAGS.gin_param)


# TODO(noam): maybe add gin-config to mtf.get_variable so we can delete
#  this stupid VariableDtype class and stop passing it all over creation.
@gin.configurable
def get_variable_dtype(
    master_dtype=tf.bfloat16,
    slice_dtype=tf.float32,
    activation_dtype=tf.float32):
  """Datatypes to use for the run.

  Args:
    master_dtype: string, datatype for checkpoints
      keep this the same between training and eval/inference
    slice_dtype: string, datatype for variables in memory
      must be tf.float32 for training
    activation_dtype: string, datatype for activations
      less memory usage if tf.bfloat16 but possible numerical issues
  Returns:
    a mtf.VariableDtype
  """
  return mtf.VariableDType(
      master_dtype=tf.as_dtype(master_dtype),
      slice_dtype=tf.as_dtype(slice_dtype),
      activation_dtype=tf.as_dtype(activation_dtype))


def inputs_vocabulary(vocabulary):
  """Inputs vocabulary.

  Args:
    vocabulary: Vocabulary or (inputs_vocabulary, targets_vocabulary) tuple.

  Returns:
    a Vocabulary
  """
  if isinstance(vocabulary, tuple):
    vocabulary = vocabulary[0]
  return vocabulary


def targets_vocabulary(vocabulary):
  """Targets vocabulary.

  Args:
    vocabulary: Vocabulary or (inputs_vocabulary, targets_vocabulary) tuple.

  Returns:
    a Vocabulary
  """
  if isinstance(vocabulary, tuple):
    vocabulary = vocabulary[1]
  return vocabulary


@gin.configurable
def separate_vocabularies(inputs=gin.REQUIRED, targets=gin.REQUIRED):
  """Gin-configurable helper function."""
  return (inputs, targets)


def build_model(model_type="bitransformer",
                input_vocab_size=gin.REQUIRED,
                output_vocab_size=gin.REQUIRED,
                layout_rules=None,
                mesh_shape=None):
  """Build a transformer model.

  Currently, four types of models are supported:

  "bitransformer": The traditional encoder-decoder architecture from
     "attention is all you need".  Requires a non-text2self dataset.

  "lm": an autoregressive language model (one layer stack).  This is similar
     to the decoder part of a bitransformer, but with no attention over an
     encoder, since there is no encoder.  Requires a text2self dataset,
     with targets, but no inputs.

  "aligned": a non-autoregressive single-stack model (like BERT).  Requires
     a non-text2self dataset with inputs and targets.  The targets are
     aligned with the inputs.

  "bi_teacher_student": a teacher-student model where both the student and
    teacher are bitransformers. Requires a non-text2self dataset.

  Args:
    model_type: a string - "bitransformer", "lm" or "aligned"
    input_vocab_size: an integer
    output_vocab_size: an integer
    layout_rules: optional - an input to mtf.convert_to_layout_rules
    mesh_shape: optional - an input to mtf.convert_to_shape
  Returns:
    a Unitransformer or Bitransformer
  """
  if model_type == "bitransformer":
    return transformer.make_bitransformer(
        input_vocab_size=input_vocab_size,
        output_vocab_size=output_vocab_size,
        mesh_shape=mesh_shape,
        layout=layout_rules)
  elif model_type == "bi_student_teacher":
    return transformer.make_bi_student_teacher(
        input_vocab_size=input_vocab_size,
        output_vocab_size=output_vocab_size,
        mesh_shape=mesh_shape,
        layout=layout_rules)
  elif model_type == "lm" or model_type == "aligned":
    return transformer.Unitransformer(
        autoregressive=model_type == "lm",
        layer_stack=transformer.make_layer_stack(),
        input_vocab_size=input_vocab_size,
        output_vocab_size=output_vocab_size,
        mesh_shape=mesh_shape,
        layout=layout_rules)
  else:
    raise ValueError("unknown model_type")


@gin.configurable
def tpu_mesh_shape(tpu_topology=gin.REQUIRED,
                   model_parallelism=gin.REQUIRED):
  """Create a mesh_shape for data-parallelism and model-parallelism on TPU.

  Example: tpu_mesh_shape("4x4", 8) -> mtf.Shape(("batch", 4), ("model", 8))
  Since there are 4x4x2=32 total cores, and we want 8-way model paralleism.

  Args:
    tpu_topology: a string - e.g. "2x2"
    model_parallelism: an integer - the number of cores per model replica
  Returns:
    a mtf.Shape
  """
  x, y = tpu_topology.split("x")
  num_cores = int(x) * int(y) * 2
  data_parallelism = num_cores // model_parallelism
  dims = []
  if data_parallelism > 1:
    dims.append(mtf.Dimension("batch", data_parallelism))
  if model_parallelism > 1:
    dims.append(mtf.Dimension("model", model_parallelism))
  return mtf.Shape(dims)


def _logical_to_physical(physical_shape, mesh_shape):
  """Mapping from logical mesh to physical TPU cores.

  This is to create the logical_to_physical mapping for SimdMeshImpl.  For 2d
  meshes, we use a tiled layout with the second logical mesh-dimension
  corresponding to position within a tile.  This tends to give better
  performance.  For non-2d meshes, we use the default mapping.

  Args:
    physical_shape: a list of integers - the physical mesh shape
    mesh_shape: a mtf.Shape
  Returns:
    a permutation of range(mesh_shape.size) or None
  """
  mesh_shape = mesh_shape.to_integer_list
  if len(mesh_shape) != 2:
    return None
  # Use "tiled" mapping of logical mesh to physical mesh.
  # The first logical-mesh dimension corresponds to which phyiscal tile
  # and the second logical-mesh dimension corresponds to location within
  # a tile.
  tile_size = mesh_shape[1] // 2  # size in chips (each with 2 cores)
  lg_tile_size = int(math.log(tile_size, 2))
  t0 = 2 ** (lg_tile_size // 2)
  t1 = tile_size // t0
  tile_shape = [t0, t1]
  tf.logging.info("Mesh shape = %s" % mesh_shape)
  tf.logging.info("Physical shape = %s" % physical_shape)
  tf.logging.info("Tile shape = %s" % tile_shape)
  _, logical_to_physical = mtf.simd_mesh_impl.tile_2d(
      physical_shape, tile_shape)
  tf.logging.info("logical_to_physical = %s" % (logical_to_physical,))
  return logical_to_physical


@gin.configurable
def tpu_estimator_model_fn(model_type,
                           transformer_model,
                           model_dir,
                           use_tpu,
                           mesh_shape,
                           layout_rules,
                           batch_size,
                           sequence_length,
                           autostack,
                           keep_checkpoint_max,
                           save_checkpoints_steps,
                           learning_rate_schedule=None,
                           optimizer=None,
                           outer_batch_size=1,
                           tpu_summaries=False,
                           predict_fn=None,
                           variable_filter=None,
                           init_checkpoint=None):
  """Create a TPUEstimator model function.

  Args:
    model_type: a string
    transformer_model: a transformer.Unitransformer or transformer.Bitransformer
    model_dir: a string
    use_tpu: a boolean
    mesh_shape: a mtf.Shape
    layout_rules: a mtf.LayoutRules
    batch_size: an integer
    sequence_length: an integer
    autostack: a boolean
    keep_checkpoint_max: an integer
    save_checkpoints_steps: an integer
    learning_rate_schedule: an optional function taking the scalar named
      argument `step` and return the scalar learning rate. Alternatively, a
      constant.
    optimizer: a class extending optimize.Optimizer, required for training
    outer_batch_size: outer batch dimension that could be used to enable the mix
      of data-parallel and model-parallel training of MoE models
    tpu_summaries: a boolean - if True, then use rewrites to make summaries work
      on TPU.  This may be slow, since it uses a host call hack.
    predict_fn: an optional function, see docs for run for more information
    variable_filter: a string, a variable will only be trained if
      this string appears in its name. If None (default), train all trainable
      variables.
    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.

  Returns:
    a function to be passed to TPUEstimator
  """
  def my_model_fn(features, labels, mode, params=None, config=None):
    """Estimator model function.

    Args:
      features: input features dictionary
      labels: ignored
      mode: a tf.estimator.ModeKeys
      params: something
      config: something

    Returns:
      something
    """
    del labels, config
    global_step = tf.train.get_global_step()
    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
      physical_shape = list(
          params["context"].device_assignment.topology.mesh_shape)
      logical_to_physical = _logical_to_physical(physical_shape, mesh_shape)
      mesh_impl = mtf.simd_mesh_impl.SimdMeshImpl(
          mesh_shape, layout_rules, mesh_devices, ctx.device_assignment,
          logical_to_physical=logical_to_physical)
    else:
      var_placer = None
      mesh_devices = [""] * mesh_shape.size
      mesh_impl = mtf.placement_mesh_impl.PlacementMeshImpl(
          mesh_shape, layout_rules, mesh_devices)

    graph = mtf.Graph()
    mesh = mtf.Mesh(graph, "my_mesh", var_placer)

    outer_batch_dim = mtf.Dimension("outer_batch", outer_batch_size)
    batch_dim = mtf.Dimension("batch", batch_size // outer_batch_size)
    length_dim = mtf.Dimension("length", sequence_length)
    feature_shape = mtf.Shape([outer_batch_dim, batch_dim, length_dim])

    mtf_features = {}
    for key, x in features.items():
      x = tf.to_int32(features[key])
      x = tf.reshape(
          x, [outer_batch_size, batch_size // outer_batch_size, sequence_length]
      )
      if not use_tpu:
        tf.logging.info("feature %s : %s" % (key, x))
        x = tf.Print(
            x, [x], "import feature %s" % key, summarize=1000, first_n=1)
      mtf_features[key] = mtf.import_fully_replicated(
          mesh, x, feature_shape, name=key)

    if mode == tf.estimator.ModeKeys.PREDICT:
      feature_shape = mtf.Shape([
          mtf.Dimension("batch", batch_size),
          mtf.Dimension("length", sequence_length)
      ])
      mtf_features = {
          k: mtf.reshape(v, feature_shape)
          for k, v in six.iteritems(mtf_features)
      }
      inputs = mtf_features["inputs"]
      if predict_fn:
        mtf_samples = predict_fn(
            model=transformer_model,
            features=mtf_features,
            variable_dtype=get_variable_dtype())
      elif isinstance(transformer_model, transformer.Unitransformer):
        mtf_samples = transformer_model.sample_autoregressive(
            inputs, variable_dtype=get_variable_dtype(),
            remove_partial_sequences=True)
      elif isinstance(transformer_model,
                      (transformer.Bitransformer, transformer.StudentTeacher)):
        mtf_samples = transformer_model.decode(
            inputs, variable_dtype=get_variable_dtype())
      else:
        raise ValueError("unrecognized class")
      mtf_samples = mtf.anonymize(mtf_samples)
      lowering = mtf.Lowering(graph, {mesh: mesh_impl}, autostack=autostack)
      outputs = lowering.export_to_tf_tensor(mtf_samples)
      predictions = {"outputs": outputs}

      # When exporting a model, we need to communicate to TF-Serving that
      # master variables need to be copied to their slave slice variables.
      # Estimator uses a Scaffold's "local_init_op" for this purpose, so we
      # augment the default "local_init_op" here.
      def scaffold_fn():
        return tf.train.Scaffold(
            local_init_op=tf.group(
                tf.train.Scaffold.default_local_init_op(),
                lowering.copy_masters_to_slices(),
                name="mtf_local_init_op"))

      return tpu_estimator.TPUEstimatorSpec(
          mode=tf.estimator.ModeKeys.PREDICT,
          predictions=predictions,
          scaffold_fn=scaffold_fn,
          prediction_hooks=[mtf.MtfRestoreHook(lowering)])

    elif mode == tf.estimator.ModeKeys.EVAL:
      raise NotImplementedError("We don't expect to use mode == eval.")

    else:
      assert mode == tf.estimator.ModeKeys.TRAIN
      num_microbatches = serialize_num_microbatches(batch_dim,
                                                    length_dim,
                                                    mesh_shape,
                                                    layout_rules)
      def model_fn(mtf_features):
        """The kind of function we need for mtf.serialize_training_step.

        Args:
          mtf_features: a dictionary
        Returns:
          a dictionary
        """
        if model_type == "lm":
          if "inputs" in mtf_features:
            mtf_features = _dynamic_text2self(mtf_features)
          _, _, length_dim = mtf_features["targets"].shape
          inputs = mtf.shift(mtf_features["targets"], offset=1,
                             dim=length_dim, wrap=False)
        else:
          inputs = mtf_features["inputs"]

        if isinstance(transformer_model, transformer.Unitransformer):
          position_kwargs = dict(
              sequence_id=mtf_features.get("targets_segmentation", None),
              position=mtf_features.get("targets_position", None),
          )
        elif isinstance(
            transformer_model,
            transformer.Bitransformer) or model_type == "bi_student_teacher":
          position_kwargs = dict(
              encoder_sequence_id=mtf_features.get("inputs_segmentation", None),
              decoder_sequence_id=mtf_features.get("targets_segmentation",
                                                   None),
              decoder_subsequence_id=mtf_features.get("targets_subsegmentation",
                                                      None),
              encoder_position=mtf_features.get("inputs_position", None),
              decoder_position=mtf_features.get("targets_position", None),
          )
        else:
          raise ValueError("unrecognized class")

        logits, loss = transformer_model.call_simple(
            inputs=inputs,
            targets=mtf_features["targets"],
            compute_loss=True,
            mode=mode,
            variable_dtype=get_variable_dtype(),
            **position_kwargs)
        if num_microbatches > 1:
          loss /= float(num_microbatches)
        del logits
        return {"loss": loss}

      if num_microbatches > 1:
        var_grads, loss_dict = mtf.serialize_training_step(
            mtf_features, model_fn, batch_dim, num_microbatches)
      else:
        loss_dict = model_fn(mtf_features)
        var_grads = mtf.gradients(
            [loss_dict["loss"]],
            [v.outputs[0] for v in graph.trainable_variables])

      loss = loss_dict["loss"]

      if callable(learning_rate_schedule):
        # the following happens on CPU since TPU can't handle summaries.
        with mtf.utils.outside_all_rewrites():
          learning_rate = learning_rate_schedule(
              step=tf.train.get_global_step())
          tf.summary.scalar("learning_rate", learning_rate)
      else:
        learning_rate = learning_rate_schedule

      filter_string = variable_filter or ""
      train_vars = graph.trainable_variables
      filtered_vars = [v for v in train_vars if filter_string in v.name]
      filtered_grads = [
          g for g, v in zip(var_grads, train_vars) if filter_string in v.name
      ]

      update_ops = optimizer(learning_rate=learning_rate).apply_grads(
          filtered_grads, filtered_vars
      )

      lowering = mtf.Lowering(graph, {mesh: mesh_impl}, autostack=autostack)

      tf_loss = lowering.export_to_tf_tensor(loss)
      tf_loss = tf.to_float(tf_loss)
      if not use_tpu:
        tf_loss = tf.Print(tf_loss, [tf_loss, tf.train.get_global_step()],
                           "step, tf_loss")

      tf_update_ops = [lowering.lowered_operation(op) for op in update_ops]
      tf_update_ops.append(tf.assign_add(global_step, 1))
      train_op = tf.group(tf_update_ops)

      if hasattr(transformer_model, "initialize"):
        with mtf.utils.outside_all_rewrites():
          transformer_model.initialize()

      with mtf.utils.outside_all_rewrites():

        if init_checkpoint:
          ckpt_vars = {v for v, _ in tf.train.list_variables(init_checkpoint)}
          global_vars = {v.op.name for v in tf.global_variables()}
          restore_vars = ckpt_vars.intersection(global_vars)
          tf.logging.info("Initializing variables from %s:", init_checkpoint)
          tf.logging.info(restore_vars)
          tf.logging.info("Variables in %s but not in graph:", init_checkpoint)
          tf.logging.info(ckpt_vars - global_vars)
          tf.logging.info("Variables in graph but not in %s:", init_checkpoint)
          tf.logging.info(global_vars - ckpt_vars)
          tf.train.init_from_checkpoint(
              init_checkpoint, {v: v for v in restore_vars}
          )

        # Copy master variables to slices. Must be called first.
        restore_hook = mtf.MtfRestoreHook(lowering)<commit_msg>Improve logging when initializing from a checkpoint.<commit_after># coding=utf-8
# Copyright 2019 The Mesh TensorFlow Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

r"""Utilities for running training and inference."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import functools
import math
import os
import re

import gin
import gin.tf

import mesh_tensorflow as mtf
from mesh_tensorflow.transformer import dataset as transformer_dataset
from mesh_tensorflow.transformer import transformer
import numpy as np
import six
import tensorflow as tf
import tensorflow_datasets as tfds

from tensorflow.contrib.tpu.python.tpu import tpu_config
from tensorflow.contrib.tpu.python.tpu import tpu_estimator

tf.flags.DEFINE_multi_string("gin_file", None,
                             "List of paths to the config files.")
tf.flags.DEFINE_multi_string(
    "gin_param", None, "Newline separated list of Gin parameter bindings.")
FLAGS = tf.flags.FLAGS

_DEFAULT_CONFIG_FILE = "./gin/defaults.gin"

# List of features used by model.
_INPUT_FEATURES = [
    "inputs", "inputs_position", "inputs_segmentation", "targets",
    "targets_position", "targets_segmentation", "targets_subsegmentation"
]


def parse_gin_defaults_and_flags():
  """Parses all default gin files and those provides via flags."""
  # Set up the default values for the configurable parameters. These values will
  # be overridden by any user provided gin files/parameters.
  gin.parse_config_file(
      os.path.join(os.path.dirname(__file__), _DEFAULT_CONFIG_FILE))
  gin.parse_config_files_and_bindings(FLAGS.gin_file, FLAGS.gin_param)


# TODO(noam): maybe add gin-config to mtf.get_variable so we can delete
#  this stupid VariableDtype class and stop passing it all over creation.
@gin.configurable
def get_variable_dtype(
    master_dtype=tf.bfloat16,
    slice_dtype=tf.float32,
    activation_dtype=tf.float32):
  """Datatypes to use for the run.

  Args:
    master_dtype: string, datatype for checkpoints
      keep this the same between training and eval/inference
    slice_dtype: string, datatype for variables in memory
      must be tf.float32 for training
    activation_dtype: string, datatype for activations
      less memory usage if tf.bfloat16 but possible numerical issues
  Returns:
    a mtf.VariableDtype
  """
  return mtf.VariableDType(
      master_dtype=tf.as_dtype(master_dtype),
      slice_dtype=tf.as_dtype(slice_dtype),
      activation_dtype=tf.as_dtype(activation_dtype))


def inputs_vocabulary(vocabulary):
  """Inputs vocabulary.

  Args:
    vocabulary: Vocabulary or (inputs_vocabulary, targets_vocabulary) tuple.

  Returns:
    a Vocabulary
  """
  if isinstance(vocabulary, tuple):
    vocabulary = vocabulary[0]
  return vocabulary


def targets_vocabulary(vocabulary):
  """Targets vocabulary.

  Args:
    vocabulary: Vocabulary or (inputs_vocabulary, targets_vocabulary) tuple.

  Returns:
    a Vocabulary
  """
  if isinstance(vocabulary, tuple):
    vocabulary = vocabulary[1]
  return vocabulary


@gin.configurable
def separate_vocabularies(inputs=gin.REQUIRED, targets=gin.REQUIRED):
  """Gin-configurable helper function."""
  return (inputs, targets)


def build_model(model_type="bitransformer",
                input_vocab_size=gin.REQUIRED,
                output_vocab_size=gin.REQUIRED,
                layout_rules=None,
                mesh_shape=None):
  """Build a transformer model.

  Currently, four types of models are supported:

  "bitransformer": The traditional encoder-decoder architecture from
     "attention is all you need".  Requires a non-text2self dataset.

  "lm": an autoregressive language model (one layer stack).  This is similar
     to the decoder part of a bitransformer, but with no attention over an
     encoder, since there is no encoder.  Requires a text2self dataset,
     with targets, but no inputs.

  "aligned": a non-autoregressive single-stack model (like BERT).  Requires
     a non-text2self dataset with inputs and targets.  The targets are
     aligned with the inputs.

  "bi_teacher_student": a teacher-student model where both the student and
    teacher are bitransformers. Requires a non-text2self dataset.

  Args:
    model_type: a string - "bitransformer", "lm" or "aligned"
    input_vocab_size: an integer
    output_vocab_size: an integer
    layout_rules: optional - an input to mtf.convert_to_layout_rules
    mesh_shape: optional - an input to mtf.convert_to_shape
  Returns:
    a Unitransformer or Bitransformer
  """
  if model_type == "bitransformer":
    return transformer.make_bitransformer(
        input_vocab_size=input_vocab_size,
        output_vocab_size=output_vocab_size,
        mesh_shape=mesh_shape,
        layout=layout_rules)
  elif model_type == "bi_student_teacher":
    return transformer.make_bi_student_teacher(
        input_vocab_size=input_vocab_size,
        output_vocab_size=output_vocab_size,
        mesh_shape=mesh_shape,
        layout=layout_rules)
  elif model_type == "lm" or model_type == "aligned":
    return transformer.Unitransformer(
        autoregressive=model_type == "lm",
        layer_stack=transformer.make_layer_stack(),
        input_vocab_size=input_vocab_size,
        output_vocab_size=output_vocab_size,
        mesh_shape=mesh_shape,
        layout=layout_rules)
  else:
    raise ValueError("unknown model_type")


@gin.configurable
def tpu_mesh_shape(tpu_topology=gin.REQUIRED,
                   model_parallelism=gin.REQUIRED):
  """Create a mesh_shape for data-parallelism and model-parallelism on TPU.

  Example: tpu_mesh_shape("4x4", 8) -> mtf.Shape(("batch", 4), ("model", 8))
  Since there are 4x4x2=32 total cores, and we want 8-way model paralleism.

  Args:
    tpu_topology: a string - e.g. "2x2"
    model_parallelism: an integer - the number of cores per model replica
  Returns:
    a mtf.Shape
  """
  x, y = tpu_topology.split("x")
  num_cores = int(x) * int(y) * 2
  data_parallelism = num_cores // model_parallelism
  dims = []
  if data_parallelism > 1:
    dims.append(mtf.Dimension("batch", data_parallelism))
  if model_parallelism > 1:
    dims.append(mtf.Dimension("model", model_parallelism))
  return mtf.Shape(dims)


def _logical_to_physical(physical_shape, mesh_shape):
  """Mapping from logical mesh to physical TPU cores.

  This is to create the logical_to_physical mapping for SimdMeshImpl.  For 2d
  meshes, we use a tiled layout with the second logical mesh-dimension
  corresponding to position within a tile.  This tends to give better
  performance.  For non-2d meshes, we use the default mapping.

  Args:
    physical_shape: a list of integers - the physical mesh shape
    mesh_shape: a mtf.Shape
  Returns:
    a permutation of range(mesh_shape.size) or None
  """
  mesh_shape = mesh_shape.to_integer_list
  if len(mesh_shape) != 2:
    return None
  # Use "tiled" mapping of logical mesh to physical mesh.
  # The first logical-mesh dimension corresponds to which phyiscal tile
  # and the second logical-mesh dimension corresponds to location within
  # a tile.
  tile_size = mesh_shape[1] // 2  # size in chips (each with 2 cores)
  lg_tile_size = int(math.log(tile_size, 2))
  t0 = 2 ** (lg_tile_size // 2)
  t1 = tile_size // t0
  tile_shape = [t0, t1]
  tf.logging.info("Mesh shape = %s" % mesh_shape)
  tf.logging.info("Physical shape = %s" % physical_shape)
  tf.logging.info("Tile shape = %s" % tile_shape)
  _, logical_to_physical = mtf.simd_mesh_impl.tile_2d(
      physical_shape, tile_shape)
  tf.logging.info("logical_to_physical = %s" % (logical_to_physical,))
  return logical_to_physical


@gin.configurable
def tpu_estimator_model_fn(model_type,
                           transformer_model,
                           model_dir,
                           use_tpu,
                           mesh_shape,
                           layout_rules,
                           batch_size,
                           sequence_length,
                           autostack,
                           keep_checkpoint_max,
                           save_checkpoints_steps,
                           learning_rate_schedule=None,
                           optimizer=None,
                           outer_batch_size=1,
                           tpu_summaries=False,
                           predict_fn=None,
                           variable_filter=None,
                           init_checkpoint=None):
  """Create a TPUEstimator model function.

  Args:
    model_type: a string
    transformer_model: a transformer.Unitransformer or transformer.Bitransformer
    model_dir: a string
    use_tpu: a boolean
    mesh_shape: a mtf.Shape
    layout_rules: a mtf.LayoutRules
    batch_size: an integer
    sequence_length: an integer
    autostack: a boolean
    keep_checkpoint_max: an integer
    save_checkpoints_steps: an integer
    learning_rate_schedule: an optional function taking the scalar named
      argument `step` and return the scalar learning rate. Alternatively, a
      constant.
    optimizer: a class extending optimize.Optimizer, required for training
    outer_batch_size: outer batch dimension that could be used to enable the mix
      of data-parallel and model-parallel training of MoE models
    tpu_summaries: a boolean - if True, then use rewrites to make summaries work
      on TPU.  This may be slow, since it uses a host call hack.
    predict_fn: an optional function, see docs for run for more information
    variable_filter: a string, a variable will only be trained if
      this string appears in its name. If None (default), train all trainable
      variables.
    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.

  Returns:
    a function to be passed to TPUEstimator
  """
  def my_model_fn(features, labels, mode, params=None, config=None):
    """Estimator model function.

    Args:
      features: input features dictionary
      labels: ignored
      mode: a tf.estimator.ModeKeys
      params: something
      config: something

    Returns:
      something
    """
    del labels, config
    global_step = tf.train.get_global_step()
    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
      physical_shape = list(
          params["context"].device_assignment.topology.mesh_shape)
      logical_to_physical = _logical_to_physical(physical_shape, mesh_shape)
      mesh_impl = mtf.simd_mesh_impl.SimdMeshImpl(
          mesh_shape, layout_rules, mesh_devices, ctx.device_assignment,
          logical_to_physical=logical_to_physical)
    else:
      var_placer = None
      mesh_devices = [""] * mesh_shape.size
      mesh_impl = mtf.placement_mesh_impl.PlacementMeshImpl(
          mesh_shape, layout_rules, mesh_devices)

    graph = mtf.Graph()
    mesh = mtf.Mesh(graph, "my_mesh", var_placer)

    outer_batch_dim = mtf.Dimension("outer_batch", outer_batch_size)
    batch_dim = mtf.Dimension("batch", batch_size // outer_batch_size)
    length_dim = mtf.Dimension("length", sequence_length)
    feature_shape = mtf.Shape([outer_batch_dim, batch_dim, length_dim])

    mtf_features = {}
    for key, x in features.items():
      x = tf.to_int32(features[key])
      x = tf.reshape(
          x, [outer_batch_size, batch_size // outer_batch_size, sequence_length]
      )
      if not use_tpu:
        tf.logging.info("feature %s : %s" % (key, x))
        x = tf.Print(
            x, [x], "import feature %s" % key, summarize=1000, first_n=1)
      mtf_features[key] = mtf.import_fully_replicated(
          mesh, x, feature_shape, name=key)

    if mode == tf.estimator.ModeKeys.PREDICT:
      feature_shape = mtf.Shape([
          mtf.Dimension("batch", batch_size),
          mtf.Dimension("length", sequence_length)
      ])
      mtf_features = {
          k: mtf.reshape(v, feature_shape)
          for k, v in six.iteritems(mtf_features)
      }
      inputs = mtf_features["inputs"]
      if predict_fn:
        mtf_samples = predict_fn(
            model=transformer_model,
            features=mtf_features,
            variable_dtype=get_variable_dtype())
      elif isinstance(transformer_model, transformer.Unitransformer):
        mtf_samples = transformer_model.sample_autoregressive(
            inputs, variable_dtype=get_variable_dtype(),
            remove_partial_sequences=True)
      elif isinstance(transformer_model,
                      (transformer.Bitransformer, transformer.StudentTeacher)):
        mtf_samples = transformer_model.decode(
            inputs, variable_dtype=get_variable_dtype())
      else:
        raise ValueError("unrecognized class")
      mtf_samples = mtf.anonymize(mtf_samples)
      lowering = mtf.Lowering(graph, {mesh: mesh_impl}, autostack=autostack)
      outputs = lowering.export_to_tf_tensor(mtf_samples)
      predictions = {"outputs": outputs}

      # When exporting a model, we need to communicate to TF-Serving that
      # master variables need to be copied to their slave slice variables.
      # Estimator uses a Scaffold's "local_init_op" for this purpose, so we
      # augment the default "local_init_op" here.
      def scaffold_fn():
        return tf.train.Scaffold(
            local_init_op=tf.group(
                tf.train.Scaffold.default_local_init_op(),
                lowering.copy_masters_to_slices(),
                name="mtf_local_init_op"))

      return tpu_estimator.TPUEstimatorSpec(
          mode=tf.estimator.ModeKeys.PREDICT,
          predictions=predictions,
          scaffold_fn=scaffold_fn,
          prediction_hooks=[mtf.MtfRestoreHook(lowering)])

    elif mode == tf.estimator.ModeKeys.EVAL:
      raise NotImplementedError("We don't expect to use mode == eval.")

    else:
      assert mode == tf.estimator.ModeKeys.TRAIN
      num_microbatches = serialize_num_microbatches(batch_dim,
                                                    length_dim,
                                                    mesh_shape,
                                                    layout_rules)
      def model_fn(mtf_features):
        """The kind of function we need for mtf.serialize_training_step.

        Args:
          mtf_features: a dictionary
        Returns:
          a dictionary
        """
        if model_type == "lm":
          if "inputs" in mtf_features:
            mtf_features = _dynamic_text2self(mtf_features)
          _, _, length_dim = mtf_features["targets"].shape
          inputs = mtf.shift(mtf_features["targets"], offset=1,
                             dim=length_dim, wrap=False)
        else:
          inputs = mtf_features["inputs"]

        if isinstance(transformer_model, transformer.Unitransformer):
          position_kwargs = dict(
              sequence_id=mtf_features.get("targets_segmentation", None),
              position=mtf_features.get("targets_position", None),
          )
        elif isinstance(
            transformer_model,
            transformer.Bitransformer) or model_type == "bi_student_teacher":
          position_kwargs = dict(
              encoder_sequence_id=mtf_features.get("inputs_segmentation", None),
              decoder_sequence_id=mtf_features.get("targets_segmentation",
                                                   None),
              decoder_subsequence_id=mtf_features.get("targets_subsegmentation",
                                                      None),
              encoder_position=mtf_features.get("inputs_position", None),
              decoder_position=mtf_features.get("targets_position", None),
          )
        else:
          raise ValueError("unrecognized class")

        logits, loss = transformer_model.call_simple(
            inputs=inputs,
            targets=mtf_features["targets"],
            compute_loss=True,
            mode=mode,
            variable_dtype=get_variable_dtype(),
            **position_kwargs)
        if num_microbatches > 1:
          loss /= float(num_microbatches)
        del logits
        return {"loss": loss}

      if num_microbatches > 1:
        var_grads, loss_dict = mtf.serialize_training_step(
            mtf_features, model_fn, batch_dim, num_microbatches)
      else:
        loss_dict = model_fn(mtf_features)
        var_grads = mtf.gradients(
            [loss_dict["loss"]],
            [v.outputs[0] for v in graph.trainable_variables])

      loss = loss_dict["loss"]

      if callable(learning_rate_schedule):
        # the following happens on CPU since TPU can't handle summaries.
        with mtf.utils.outside_all_rewrites():
          learning_rate = learning_rate_schedule(
              step=tf.train.get_global_step())
          tf.summary.scalar("learning_rate", learning_rate)
      else:
        learning_rate = learning_rate_schedule

      filter_string = variable_filter or ""
      train_vars = graph.trainable_variables
      filtered_vars = [v for v in train_vars if filter_string in v.name]
      filtered_grads = [
          g for g, v in zip(var_grads, train_vars) if filter_string in v.name
      ]

      update_ops = optimizer(learning_rate=learning_rate).apply_grads(
          filtered_grads, filtered_vars
      )

      lowering = mtf.Lowering(graph, {mesh: mesh_impl}, autostack=autostack)

      tf_loss = lowering.export_to_tf_tensor(loss)
      tf_loss = tf.to_float(tf_loss)
      if not use_tpu:
        tf_loss = tf.Print(tf_loss, [tf_loss, tf.train.get_global_step()],
                           "step, tf_loss")

      tf_update_ops = [lowering.lowered_operation(op) for op in update_ops]
      tf_update_ops.append(tf.assign_add(global_step, 1))
      train_op = tf.group(tf_update_ops)

      if hasattr(transformer_model, "initialize"):
        with mtf.utils.outside_all_rewrites():
          transformer_model.initialize()

      with mtf.utils.outside_all_rewrites():

        if init_checkpoint:
          ckpt_vars = {v for v, _ in tf.train.list_variables(init_checkpoint)}
          global_vars = {v.op.name for v in tf.global_variables()}
          restore_vars = ckpt_vars.intersection(global_vars)
          tf.logging.info("Initializing variables from %s:", init_checkpoint)
          tf.logging.info("\n".join(sorted(restore_vars)))
          tf.logging.info("Variables in %s but not in graph:", init_checkpoint)
          tf.logging.info("\n".join(sorted(ckpt_vars - global_vars)))
          tf.logging.info("Variables in graph but not in %s:", init_checkpoint)
          tf.logging.info("\n".join(sorted(global_vars - ckpt_vars)))
          tf.train.init_from_checkpoint(
              init_checkpoint, {v: v for v in restore_vars}
          )

        # Copy master variables to slices. Must be called first.
        restore_hook = mtf.MtfRestoreHook(lowering)
{'fraction_non_alphanumeric': 0.07905148237179487, 'fraction_numerical': 0.004407051282051282, 'mean_word_length': 2.766220294228593, 'pattern_counts': {'":': 22, '<': 3, '<?xml version=': 0, '>': 13, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 24, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '957386', 'n_tokens_mistral': 12088, 'n_tokens_neox': 11696, 'n_words': 3041}