code: def expert_to_batch_indices(self):
      
      return tf.split(
          self._batch_index, self._part_sizes_tensor, 0, num=self._num_experts)
docs: Batch indices corresponding to the examples in the per-expert `Tensor`s.
  
      Returns:
        a list of `num_experts` one-dimensional `Tensor`s with type `tf.int64`
            and shapes `[expert_batch_size_i]`