def model_fn_sharded(self, sharded_features):
    """Estimator model_fn sharded along batch dimension.

    Args:
      sharded_features: {str: [Tensor]}. Features sharded along batch dimension.
        Each list is the same length (== number of shards).

    Returns:
      sharded_logits: [Tensor]. Logits for each shard of examples.
      losses: {str: 0-D Tensor}. Loss averaged across shards.
    """
    dp = self._data_parallelism

    # [{str: Tensor}]. Transpose of 'sharded_features'.
    datashard_to_features = self._to_features_per_datashard(sharded_features)
    sharded_logits, sharded_losses, sharded_monitors, _ = (
        dp(self.model_fn, datashard_to_features))
    sharded_logits, sharded_losses = dp(
        self.maybe_scheduled_sampling,
        datashard_to_features, sharded_logits, sharded_losses)
    if isinstance(sharded_logits[0], dict):
      temp_dict = {k: [] for k, _ in six.iteritems(sharded_logits[0])}
      for k, _ in six.iteritems(sharded_logits[0]):
        for l in sharded_logits:
          temp_dict[k].append(l[k])
      sharded_logits = temp_dict
    losses = t2t_model.average_sharded_losses(sharded_losses)
    monitor = {}
    for key in list(sharded_monitors[0].keys()):
      monitor[key] = (
          tf.add_n([m[key] for m in sharded_monitors]) / len(sharded_monitors))
    ops.save_summary(monitor, "monitor")

    return sharded_logits, losses
Estimator model_fn sharded along batch dimension.

    Args:
      sharded_features: {str: [Tensor]}. Features sharded along batch dimension.
        Each list is the same length (== number of shards).

    Returns:
      sharded_logits: [Tensor]. Logits for each shard of examples.
      losses: {str: 0-D Tensor}. Loss averaged across shards.
    
tensor2tensor/models/research/transformer_vae_flow_prior.py