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`.
class_docstr_tok: ["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.
  
  Args:
    num_experts: an integer.
    gates: a `Tensor` of shape `[batch_size, num_experts]`.
  
  Returns:
    a SparseDispatcher
docstr_tok: ["Create","a","SparseDispatcher",".","Args",":","num_experts",":","an","integer",".","gates",":","a","`","Tensor","`","of","shape","`","[","batch_size",",","num_experts","]","`",".","Returns",":","a","SparseDispatcher"]
code: def __init__(self, num_experts, gates):
      """Create a SparseDispatcher.
  
      Args:
        num_experts: an integer.
        gates: a `Tensor` of shape `[batch_size, num_experts]`.
  
      Returns:
        a SparseDispatcher
      """
      self._gates = gates
      self._num_experts = num_experts
  
      where = tf.to_int32(tf.where(tf.transpose(gates) > 0))
      self._expert_index, self._batch_index = tf.unstack(
          where, num=2, axis=1
      )
      self._part_sizes_tensor = tf.reduce_sum(
          tf.to_int32(gates > 0), [0]
      )
      self._nonzero_gates = tf.gather(
          tf.reshape(self._gates, [-1]),
          self._batch_index * num_experts + self._expert_index,
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import functools
  import math
  import six
  from six.moves import range
  from six.moves import zip
  from tensor2tensor.layers import common_layers
  from tensor2tensor.layers.vq_discrete import DiscreteBottleneck
  import tensorflow.compat.v1
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)