code: def dispatch(self, inp):
      
      inp = tf.gather(inp, self._batch_index)
      return tf.split(inp, self._part_sizes_tensor, 0, num=self._num_experts)
docs: Create one input Tensor for each expert.
  
      The `Tensor` for a expert `i` contains the slices of `inp` corresponding
      to the batch elements `b` where `gates[b, i] > 0`.
  
      Args:
        inp: a `Tensor` of shape "[batch_size, <extra_input_dims>]`
      Returns:
        a list of `num_experts` `Tensor`s with shapes
          `[expert_batch_size_i, <extra_input_dims>]`.