code: def dispatch(self, inp):
      
      dispatched = self._dp(lambda a, b: a.dispatch(b), self._dispatchers, inp)
      ret = self._ep(tf.concat, transpose_list_of_lists(dispatched), 0)
      if ret[0].dtype == tf.float32:
        # see comments on common_layers.convert_gradient_to_tensor
        ret = self._ep(common_layers.convert_gradient_to_tensor, ret)
      return ret
docs: Create one input Tensor for each expert.
  
      Args:
        inp: a list of length num_datashards `Tensor`s with shapes
          `[batch_size[d], <extra_input_dims>]`.
      Returns:
        a list of `num_experts` `Tensor`s with shapes
          `[num_examples[i], <extra_input_dims>]`.