<SYSTEM_TASK:>
Solve the following problem using Python, implementing the functions described below, one line at a time
<END_TASK>
<USER_TASK:>
Description:
 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)
    if self.use_body_sharded():
      if  self.hparams.scheduled_sampling_prob > 0.0:
        raise NotImplementedError(
            "Scheduled sampling for non-sharded body only.")

      # MoE models override body_sharded
      transformed_features = dp(self.bottom, datashard_to_features)
      body_out = self.body_sharded(
          self._to_single_features_dict(transformed_features))
      body_out, losses = self._normalize_body_output(body_out)
      if "training" in losses:
        log_info("Skipping T2TModel top and loss because training loss "
                 "returned from body")
        sharded_logits = body_out
      else:
        if isinstance(body_out, dict):
          sharded_logits = collections.OrderedDict()
          sharded_losses = collections.OrderedDict()
          for k, v in sorted(six.iteritems(body_out)):
            sharded_logits[k] = dp(self.top, v, datashard_to_features)
            sharded_losses[k] = dp(self.loss, sharded_logits[k],
                                   datashard_to_features)
          training_loss_dict = average_sharded_losses([({
              "training": l
          } for l in loss) for loss in sharded_losses.values()])
          losses.update(training_loss_dict)
        else:
          sharded_logits = dp(self.top, body_out, datashard_to_features)
          sharded_losses = dp(self.loss, sharded_logits, datashard_to_features)
          if isinstance(sharded_losses, tuple):
            nums, dens = sharded_losses
            sharded_losses = zip(nums, dens)
          training_loss_dict = average_sharded_losses([{
              "training": loss
          } for loss in sharded_losses])
          losses.update(training_loss_dict)
    else:
      sharded_logits, sharded_losses = 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 = average_sharded_losses(sharded_losses)

    return sharded_logits, losses