code: def dispatch(self, inp):
      
      inp = tf.reshape(inp, [self._batch * self._length, -1])
      # [batch, num_experts, expert_capacity, depth]
      ret = tf.gather(inp, self._flat_indices)
      return ret
docs: Send the inputs to the experts.
  
      Args:
        inp: a `Tensor` of shape "[batch, length, depth]`
      Returns:
        a tensor with shape [batch, num_experts, expert_capacity, depth]