class_docstr: Helper for implementing a mixture of experts.
  The purpose of this class is to create input minibatches for the
  experts and to combine the results of the experts to form a unified
  output tensor.
  There are two functions:
  dispatch - take an input Tensor and create input Tensors for each expert.
  combine - take output Tensors from each expert and form a combined output
    Tensor.  Outputs from different experts for the same batch element are
    summed together, weighted by the provided "gates".
  The class is initialized with a "gates" Tensor, which specifies which
  batch elements go to which experts, and the weights to use when combining
  the outputs.  Batch element b is sent to expert e iff gates[b, e] != 0.
  The inputs and outputs are all two-dimensional [batch, depth].
  Caller is responsible for collapsing additional dimensions prior to
  calling this class and reshaping the output to the original shape.
  See common_layers.reshape_like().
  Example use:
  gates: a float32 `Tensor` with shape `[batch_size, num_experts]`
  inputs: a float32 `Tensor` with shape `[batch_size, input_size]`
  experts: a list of length `num_experts` containing sub-networks.
  dispatcher = SparseDispatcher(num_experts, gates)
  expert_inputs = dispatcher.dispatch(inputs)
  expert_outputs = [experts[i](expert_inputs[i]) for i in range(num_experts)]
  outputs = dispatcher.combine(expert_outputs)
  The preceding code sets the output for a particular example b to:
  output[b] = Sum_i(gates[b, i] * experts[i](inputs[b]))
  This class takes advantage of sparsity in the gate matrix by including in the
  `Tensor`s for expert i only the batch elements for which `gates[b, i] > 0`.
docstr: Create a SparseDispatcher.
docstr_tok: ["Create","a","SparseDispatcher","."]
code: def __init__(self, num_experts, gates):
      """Create a SparseDispatcher."""
  
      self._gates = gates
      self._num_experts = num_experts
      # sort experts
      sorted_experts, index_sorted_experts = torch.nonzero(gates).sort(
          0
      )
      # drop indices
      _, self._expert_index = sorted_experts.split(1, dim=1)
      # get according batch index for each expert
      self._batch_index = torch.nonzero(gates)[
          index_sorted_experts[:, 1], 0
      ]
      # calculate num samples that each expert gets
      self._part_sizes = (gates > 0).sum(0).tolist()
      # expand gates to match with self._batch_index
      gates_exp = gates[self._batch_index.flatten()]
      self._nonzero_gates = torch.gather(
          gates_exp, 1, self._expert_index
      )
  
imports: import torch
  import torch.nn
  from torch.distributions.normal import Normal
  import numpy
label_desc: Use image node_id 1 to create a new SparseDispatcher object from inherited base classes: object with example: obj = SparseDispatcher(num_experts, gates)