Python tensorflow.compat.v1.gather() Examples

def testGather(self):
    gather_index = [5, 6, 7, 8, 9, 0, 1, 2, 3, 4]
    with arg_scope(self._batch_norm_scope()):
      inputs = tf.zeros([2, 4, 4, 3])
      c1 = layers.conv2d(inputs, num_outputs=10, kernel_size=3, scope='conv1')
      gather = tf.gather(c1, gather_index, axis=3)

    manager = orm.OpRegularizerManager(
        [gather.op], self._default_op_handler_dict)

    c1_reg = manager.get_regularizer(_get_op('conv1/Conv2D'))
    gather_reg = manager.get_regularizer(_get_op('GatherV2'))

    # Check regularizer indices.
    self.assertAllEqual(list(range(10)), c1_reg.regularization_vector)
    # This fails due to gather not being supported.  Once gather is supported,
    # this test can be enabled to verify that the regularization vector is
    # gathered in the same ordering as the tensor.
    # self.assertAllEqual(
    #     gather_index, gather_reg.regularization_vector)

    # This test shows that gather is not supported.  The regularization vector
    # has the same initial ordering after the gather op scrambled the
    # channels.  Remove this once gather is supported.
    self.assertAllEqual(list(range(10)), gather_reg.regularization_vector) 

def shuffle_layer(inputs, shuffle_fn=rol):
  """Shuffles the elements according to bitwise left or right rotation.

  Args:
    inputs: Tensor input from previous layer
    shuffle_fn: Shift function rol or ror

  Returns:
    tf.Tensor: Inputs shifted according to shuffle_fn
  """

  length = tf.shape(inputs)[1]
  n_bits = tf.log(tf.cast(length - 1, tf.float32)) / tf.log(2.0)
  n_bits = tf.cast(n_bits, tf.int32) + 1

  indices = tf.range(0, length)
  rev_indices = shuffle_fn(indices, n_bits)
  return tf.gather(inputs, rev_indices, axis=1) 

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) 

def simulate(self, action):
    reward, done = self._batch_env.simulate(action)
    with tf.control_dependencies([reward, done]):
      new_observ = tf.expand_dims(self._batch_env.observ, axis=1)

      # If we shouldn't stack, i.e. self.history == 1, then just assign
      # new_observ to self._observ and return from here.
      if self.history == 1:
        with tf.control_dependencies([self._observ.assign(new_observ)]):
          return tf.identity(reward), tf.identity(done)

      # If we should stack, then do the required work.
      old_observ = tf.gather(
          self._observ.read_value(),
          list(range(1, self.history)),
          axis=1)
      with tf.control_dependencies([new_observ, old_observ]):
        with tf.control_dependencies([self._observ.assign(
            tf.concat([old_observ, new_observ], axis=1))]):
          return tf.identity(reward), tf.identity(done) 

def post_attention(self, token, x):
    """Called after self-attention. The memory can be updated here.

    Args:
      token: Data returned by pre_attention, which can be used to carry over
        state related to the current memory operation.
      x: a Tensor of data after self-attention and feed-forward
    Returns:
      a (possibly modified) version of the input x
    """
    with tf.variable_scope(self.name + "/post_attention", reuse=tf.AUTO_REUSE):
      depth = common_layers.shape_list(x)[-1]
      actual_batch_size = common_layers.shape_list(x)[0]
      memory_output = tf.gather(token["retrieved_mem"],
                                tf.range(actual_batch_size))
      output = tf.add(tf.layers.dense(x, depth, use_bias=False),
                      tf.layers.dense(memory_output, depth))
      with tf.control_dependencies([output]):
        with tf.control_dependencies([
            self.write(token["x"], token["access_logits"])]):
          return tf.identity(output) 

def _compute_auxiliary_structure(self, contents_and_mask):
    """Compute segment and position metadata."""
    contents = contents_and_mask[:, :self._num_sequences]
    start_mask = tf.cast(contents_and_mask[:, self._num_sequences:],
                         dtype=INDEX_DTYPE)

    segment = tf.cumsum(start_mask, axis=0)
    uniform_count = tf.ones_like(segment[:, 0])
    position = []
    for i in range(self._num_sequences):
      segment_slice = segment[:, i]
      counts = tf.math.segment_sum(uniform_count, segment[:, i])
      position.append(tf.range(self._packed_length) -  tf.cumsum(
          tf.gather(counts, segment_slice - 1) * start_mask[:, i]))
    position = tf.concat([i[:, tf.newaxis] for i in position], axis=1)

    # Correct for padding tokens.
    pad_mask = tf.cast(tf.not_equal(contents, 0), dtype=INDEX_DTYPE)
    segment *= pad_mask
    position *= pad_mask

    return segment, position 

def _build_classification_specification(label, num_classes, collapse):
  """Returns a LinearSpecification for adversarial classification."""
  # Pre-construct the specifications of the different classes.
  eye = np.eye(num_classes - 1)
  specifications = []
  for i in range(num_classes):
    specifications.append(np.concatenate(
        [eye[:, :i], -np.ones((num_classes - 1, 1)), eye[:, i:]], axis=1))
  specifications = np.array(specifications, dtype=np.float32)
  specifications = tf.constant(specifications)
  # We can then use gather.
  c = tf.gather(specifications, label)
  # By construction all specifications are relevant.
  d = tf.zeros(shape=(tf.shape(label)[0], num_classes - 1))
  return ibp.LinearSpecification(c, d, prune_irrelevant=False,
                                 collapse=collapse) 

def targeted_objective(final_w, final_b, labels):
  """Determines final layer weights for attacks targeting each class.

  Args:
    final_w: 2D tensor of shape (last_hidden_layer_size, num_classes)
      containing the weights for the final linear layer.
    final_b: 1D tensor of shape (num_classes) containing the biases for the
      final hidden layer.
    labels: 1D integer tensor of shape (batch_size) of labels for each
      input example.

  Returns:
    obj_w: Tensor of shape (num_classes, batch_size, last_hidden_layer_size)
      containing weights (to use in place of final linear layer weights)
      for targeted attacks.
    obj_b: Tensor of shape (num_classes, batch_size) containing bias
      (to use in place of final linear layer biases) for targeted attacks.
  """
  # Elide objective with final linear layer.
  final_wt = tf.transpose(final_w)
  obj_w = tf.expand_dims(final_wt, axis=1) - tf.gather(final_wt, labels, axis=0)
  obj_b = tf.expand_dims(final_b, axis=1) - tf.gather(final_b, labels, axis=0)
  return obj_w, obj_b 

def slice(self, tf_tensor, tensor_shape):
    """"Slice out the corresponding part of tensor given the pnum variable."""
    tensor_layout = self.tensor_layout(tensor_shape)

    if tensor_layout.is_fully_replicated:
      return self.LaidOutTensor([tf_tensor])
    else:
      slice_shape = self.slice_shape(tensor_shape)
      slice_begins = [
          self.slice_begin(tensor_shape, pnum) for pnum in xrange(self.size)
      ]
      slice_begins_tensor = tf.stack(slice_begins)
      # slice on source device
      selected_slice_begin = tf.gather(slice_begins_tensor, self.pnum_tensor)
      return self.LaidOutTensor(
          [tf.slice(tf_tensor, selected_slice_begin, slice_shape)]) 

def entity_emb(entity_ind, entity_word_ids, entity_word_masks, word_emb_table,
               word_weights):
  """Get BOW embeddings for entities."""
  # [NNZ, max_entity_len]
  batch_entity_word_ids = tf.gather(entity_word_ids, entity_ind)
  batch_entity_word_masks = tf.gather(entity_word_masks, entity_ind)
  # [NNZ, max_entity_len, dim]
  batch_entity_word_emb = tf.gather(word_emb_table, batch_entity_word_ids)
  # [NNZ, max_entity_len, 1]
  batch_entity_word_weights = tf.gather(word_weights, batch_entity_word_ids)
  # [NNZ, dim]
  batch_entity_emb = tf.reduce_sum(
      batch_entity_word_emb * batch_entity_word_weights *
      tf.expand_dims(batch_entity_word_masks, 2),
      axis=1)
  return batch_entity_emb 

def compute_probs_tf(slate, scores_tf, score_no_click_tf):
  """Computes the selection probability and returns selected index.

  This assumes scores are normalizable, e.g., scores cannot be negative.

  Args:
    slate: a list of integers that represents the video slate.
    scores_tf: a float tensor that stores the scores of all documents.
    score_no_click_tf: a float tensor that represents the score for the action
      of picking no document.

  Returns:
    A float tensor that represents the probabilities of selecting each document
      in the slate.
  """
  all_scores = tf.concat([
      tf.gather(scores_tf, slate),
      tf.reshape(score_no_click_tf, (1, 1))
  ], axis=0)  # pyformat: disable
  all_probs = all_scores / tf.reduce_sum(input_tensor=all_scores)
  return all_probs[:-1] 

def permute_noise_tokens(tokens, noise_mask, unused_vocabulary):
  """Permute the noise tokens, keeping the non-noise tokens where they are.

  Args:
    tokens: a 1d integer Tensor
    noise_mask: a boolean Tensor with the same shape as tokens
    unused_vocabulary: a vocabulary.Vocabulary
  Returns:
    a Tensor with the same shape and dtype as tokens
  """
  masked_only = tf.boolean_mask(tokens, noise_mask)
  permuted = tf.random.shuffle(masked_only)
  # pad to avoid errors when it has size 0
  permuted = tf.pad(permuted, [[0, 1]])
  indices = tf.cumsum(tf.cast(noise_mask, tf.int32), exclusive=True)
  return tf.where_v2(noise_mask,
                     tf.gather(permuted, indices),
                     tokens) 

def _expansion_box_field_labels(self,
                                  object_classes,
                                  object_field,
                                  copy_class_id=False):
    """Expand the labels of a specific object field according to the hierarchy.

    Args:
      object_classes: Int64 tensor with the class id for each element in
        object_field.
      object_field: Tensor to be expanded.
      copy_class_id: Boolean to choose whether to use class id values in the
        output tensor instead of replicating the original values.

    Returns:
      A tensor with the result of expanding object_field.
    """
    expanded_indices = tf.gather(
        self._ancestors_lut, object_classes - _LABEL_OFFSET, axis=0)
    if copy_class_id:
      new_object_field = tf.where(expanded_indices > 0)[:, 1] + _LABEL_OFFSET
    else:
      new_object_field = tf.repeat(
          object_field, tf.reduce_sum(expanded_indices, axis=1), axis=0)
    return new_object_field 

def matmul_gather_on_zeroth_axis(params, indices, scope=None):
  """Matrix multiplication based implementation of tf.gather on zeroth axis.

  TODO(rathodv, jonathanhuang): enable sparse matmul option.

  Args:
    params: A float32 Tensor. The tensor from which to gather values.
      Must be at least rank 1.
    indices: A Tensor. Must be one of the following types: int32, int64.
      Must be in range [0, params.shape[0])
    scope: A name for the operation (optional).

  Returns:
    A Tensor. Has the same type as params. Values from params gathered
    from indices given by indices, with shape indices.shape + params.shape[1:].
  """
  with tf.name_scope(scope, 'MatMulGather'):
    params_shape = shape_utils.combined_static_and_dynamic_shape(params)
    indices_shape = shape_utils.combined_static_and_dynamic_shape(indices)
    params2d = tf.reshape(params, [params_shape[0], -1])
    indicator_matrix = tf.one_hot(indices, params_shape[0])
    gathered_result_flattened = tf.matmul(indicator_matrix, params2d)
    return tf.reshape(gathered_result_flattened,
                      tf.stack(indices_shape + params_shape[1:])) 

def clip_tensor(t, length):
  """Clips the input tensor along the first dimension up to the length.

  Args:
    t: the input tensor, assuming the rank is at least 1.
    length: a tensor of shape [1]  or an integer, indicating the first dimension
      of the input tensor t after clipping, assuming length <= t.shape[0].

  Returns:
    clipped_t: the clipped tensor, whose first dimension is length. If the
      length is an integer, the first dimension of clipped_t is set to length
      statically.
  """
  clipped_t = tf.gather(t, tf.range(length))
  if not _is_tensor(length):
    clipped_t = _set_dim_0(clipped_t, length)
  return clipped_t 

def _gather_instance_masks(self, instance_masks, classes):
    """Gathers the masks that correspond to classes.

    Args:
      instance_masks: A 4-D float32 tensor with shape
        [K, num_classes, mask_height, mask_width].
      classes: A 2-D int32 tensor with shape [batch_size, max_detection].

    Returns:
      masks: a 3-D float32 tensor with shape [K, mask_height, mask_width].
    """
    _, num_classes, height, width = instance_masks.get_shape().as_list()
    k = tf.shape(instance_masks)[0]
    instance_masks = tf.reshape(instance_masks, [-1, height, width])
    classes = tf.cast(tf.reshape(classes, [-1]), dtype=tf.int32)
    gather_idx = tf.range(k) * num_classes + classes
    return tf.gather(instance_masks, gather_idx) 

def equalize(image):
  """Implements Equalize function from PIL using TF ops."""
  def scale_channel(im, c):
    """Scale the data in the channel to implement equalize."""
    im = tf.cast(im[:, :, c], tf.int32)
    # Compute the histogram of the image channel.
    histo = tf.histogram_fixed_width(im, [0, 255], nbins=256)

    # For the purposes of computing the step, filter out the nonzeros.
    nonzero = tf.where(tf.not_equal(histo, 0))
    nonzero_histo = tf.reshape(tf.gather(histo, nonzero), [-1])
    step = (tf.reduce_sum(nonzero_histo) - nonzero_histo[-1]) // 255

    def build_lut(histo, step):
      # Compute the cumulative sum, shifting by step // 2
      # and then normalization by step.
      lut = (tf.cumsum(histo) + (step // 2)) // step
      # Shift lut, prepending with 0.
      lut = tf.concat([[0], lut[:-1]], 0)
      # Clip the counts to be in range.  This is done
      # in the C code for image.point.
      return tf.clip_by_value(lut, 0, 255)

    # If step is zero, return the original image.  Otherwise, build
    # lut from the full histogram and step and then index from it.
    result = tf.cond(tf.equal(step, 0),
                     lambda: im,
                     lambda: tf.gather(build_lut(histo, step), im))

    return tf.cast(result, tf.uint8)

  # Assumes RGB for now.  Scales each channel independently
  # and then stacks the result.
  s1 = scale_channel(image, 0)
  s2 = scale_channel(image, 1)
  s3 = scale_channel(image, 2)
  image = tf.stack([s1, s2, s3], 2)
  return image 

def _gather_valid_indices(tensor, indices, padding_value=0.0):
  """Gather values for valid indices.

  TODO(rathodv): We can't use ops.gather_with_padding_values due to cyclic
  dependency. Start using it after migrating all users of spatial ops to import
  this module directly rather than util/ops.py

  Args:
    tensor: A tensor to gather valid values from.
    indices: A 1-D int32 tensor containing indices along axis 0 of `tensor`.
      Invalid indices must be marked with -1.
    padding_value: Value to return for invalid indices.

  Returns:
    A tensor sliced based on indices. For indices that are equal to -1, returns
    rows of padding value.
  """
  padded_tensor = tf.concat(
      [
          padding_value *
          tf.ones([1, tf.shape(tensor)[-1]], dtype=tensor.dtype), tensor
      ],
      axis=0,
  )
  # tf.concat gradient op uses tf.where(condition) (which is not
  # supported on TPU) when the inputs to it are tf.IndexedSlices instead of
  # tf.Tensor. Since gradient op for tf.gather returns tf.IndexedSlices,
  # we add a dummy op inbetween tf.concat and tf.gather to ensure tf.concat
  # gradient function gets tf.Tensor inputs and not tf.IndexedSlices.
  padded_tensor *= 1.0
  return tf.gather(padded_tensor, indices + 1) 

def gather_indexes(sequence_tensor, positions):
  """Gathers the vectors at the specific positions over a minibatch."""
  sequence_shape = modeling.get_shape_list(sequence_tensor, expected_rank=3)
  batch_size = sequence_shape[0]
  seq_length = sequence_shape[1]
  width = sequence_shape[2]

  flat_offsets = tf.reshape(
      tf.range(0, batch_size, dtype=tf.int32) * seq_length, [-1, 1])
  flat_positions = tf.reshape(positions + flat_offsets, [-1])
  flat_sequence_tensor = tf.reshape(sequence_tensor,
                                    [batch_size * seq_length, width])
  output_tensor = tf.gather(flat_sequence_tensor, flat_positions)
  return output_tensor 

def noise_token_to_gathered_token(tokens, noise_mask, unused_vocabulary):
  """Replace each noise token with a random token from the sequence.

  Args:
    tokens: a 1d integer Tensor
    noise_mask: a boolean Tensor with the same shape as tokens
    unused_vocabulary: a vocabulary.Vocabulary
  Returns:
    a Tensor with the same shape and dtype as tokens
  """
  indices = tf.random_uniform(
      shape=tf.shape(tokens), maxval=tf.size(tokens), dtype=tf.int32)
  return tf.where_v2(noise_mask,
                     tf.gather(tokens, indices),
                     tokens) 

def compute_target_greedy_q(reward, gamma, next_actions, next_q_values,
                            next_states, terminals):
  """Computes the optimal target Q value with the adaptive greedy algorithm.

  This algorithm corresponds to the method "GT" in
  Ie et al. https://arxiv.org/abs/1905.12767..

  Args:
    reward: [batch_size] tensor, the immediate reward.
    gamma: float, discount factor with the usual RL meaning.
    next_actions: [batch_size, slate_size] tensor, the next slate.
    next_q_values: [batch_size, num_of_documents] tensor, the q values of the
      documents in the next step.
    next_states: [batch_size, 1 + num_of_documents] tensor, the features for the
      user and the docuemnts in the next step.
    terminals: [batch_size] tensor, indicating if this is a terminal step.

  Returns:
    [batch_size] tensor, the target q values.
  """
  slate_size = next_actions.get_shape().as_list()[1]
  stack_number = -1
  user_obs = next_states[:, 0, :, stack_number]
  doc_obs = next_states[:, 1:, :, stack_number]

  batch_size = next_q_values.get_shape().as_list()[0]
  next_greedy_q_list = []
  for i in range(batch_size):
    s, s_no_click = score_documents_tf(user_obs[i], doc_obs[i])
    q = next_q_values[i]

    slate = select_slate_greedy(slate_size, s_no_click, s, q)
    p_selected = compute_probs_tf(slate, s, s_no_click)
    q_selected = tf.gather(q, slate)
    next_greedy_q_list.append(
        tf.reduce_sum(input_tensor=p_selected * q_selected))

  next_greedy_q_values = tf.stack(next_greedy_q_list)

  return reward + gamma * next_greedy_q_values * (
      1. - tf.cast(terminals, tf.float32)) 

def compute_target_sarsa(reward, gamma, next_actions, next_q_values,
                         next_states, terminals):
  """Computes the SARSA target Q value.

  Args:
    reward: [batch_size] tensor, the immediate reward.
    gamma: float, discount factor with the usual RL meaning.
    next_actions: [batch_size, slate_size] tensor, the next slate.
    next_q_values: [batch_size, num_of_documents] tensor, the q values of the
      documents in the next step.
    next_states: [batch_size, 1 + num_of_documents] tensor, the features for the
      user and the docuemnts in the next step.
    terminals: [batch_size] tensor, indicating if this is a terminal step.

  Returns:
    [batch_size] tensor, the target q values.
  """
  stack_number = -1
  user_obs = next_states[:, 0, :, stack_number]
  doc_obs = next_states[:, 1:, :, stack_number]

  batch_size = next_q_values.get_shape().as_list()[0]
  next_sarsa_q_list = []
  for i in range(batch_size):
    s, s_no_click = score_documents_tf(user_obs[i], doc_obs[i])
    q = next_q_values[i]

    slate = tf.expand_dims(next_actions[i], 1)
    p_selected = compute_probs_tf(slate, s, s_no_click)
    q_selected = tf.gather(q, slate)
    next_sarsa_q_list.append(
        tf.reduce_sum(input_tensor=p_selected * q_selected))

  next_sarsa_q_values = tf.stack(next_sarsa_q_list)

  return reward + gamma * next_sarsa_q_values * (1. -
                                                 tf.cast(terminals, tf.float32))


# The top-K and greedy algorithms rely on the fact that the
# probs[slate] = normalize(scores[slate], score_no_click) 

def select_slate_optimal(slate_size, s_no_click, s, q):
  """Selects the slate using exhaustive search.

  This algorithm corresponds to the method "OS" in
  Ie et al. https://arxiv.org/abs/1905.12767.

  Args:
    slate_size: int, the size of the recommendation slate.
    s_no_click: float tensor, the score for not clicking any document.
    s: [num_of_documents] tensor, the scores for clicking documents.
    q: [num_of_documents] tensor, the predicted q values for documents.

  Returns:
    [slate_size] tensor, the selected slate.
  """

  num_candidates = s.shape.as_list()[0]

  # Obtain all possible slates given current docs in the candidate set.
  mesh_args = [list(range(num_candidates))] * slate_size
  slates = tf.stack(tf.meshgrid(*mesh_args), axis=-1)
  slates = tf.reshape(slates, shape=(-1, slate_size))

  # Filter slates that include duplicates to ensure each document is picked
  # at most once.
  unique_mask = tf.map_fn(
      lambda x: tf.equal(tf.size(input=x), tf.size(input=tf.unique(x)[0])),
      slates,
      dtype=tf.bool)
  slates = tf.boolean_mask(tensor=slates, mask=unique_mask)

  slate_q_values = tf.gather(s * q, slates)
  slate_scores = tf.gather(s, slates)
  slate_normalizer = tf.reduce_sum(
      input_tensor=slate_scores, axis=1) + s_no_click

  slate_q_values = slate_q_values / tf.expand_dims(slate_normalizer, 1)
  slate_sum_q_values = tf.reduce_sum(input_tensor=slate_q_values, axis=1)
  max_q_slate_index = tf.argmax(input=slate_sum_q_values)
  return tf.gather(slates, max_q_slate_index, axis=0) 

def select_slate_greedy(slate_size, s_no_click, s, q):
  """Selects the slate using the adaptive greedy algorithm.

  This algorithm corresponds to the method "GS" in
  Ie et al. https://arxiv.org/abs/1905.12767.

  Args:
    slate_size: int, the size of the recommendation slate.
    s_no_click: float tensor, the score for not clicking any document.
    s: [num_of_documents] tensor, the scores for clicking documents.
    q: [num_of_documents] tensor, the predicted q values for documents.

  Returns:
    [slate_size] tensor, the selected slate.
  """

  def argmax(v, mask):
    return tf.argmax(
        input=(v - tf.reduce_min(input_tensor=v) + 1) * mask, axis=0)

  numerator = tf.constant(0.)
  denominator = tf.constant(0.) + s_no_click
  mask = tf.ones(tf.shape(input=q)[0])

  def set_element(v, i, x):
    mask = tf.one_hot(i, tf.shape(input=v)[0])
    v_new = tf.ones_like(v) * x
    return tf.where(tf.equal(mask, 1), v_new, v)

  for _ in range(slate_size):
    k = argmax((numerator + s * q) / (denominator + s), mask)
    mask = set_element(mask, k, 0)
    numerator = numerator + tf.gather(s * q, k)
    denominator = denominator + tf.gather(s, k)

  output_slate = tf.where(tf.equal(mask, 0))
  return output_slate 

def batch_multiply(sp_tensor, dense_vec):
  """Batch multiply a vector with a sparse tensor."""
  batch_indices = sp_tensor.indices[:, 0]
  batch_vec = tf.gather(dense_vec, batch_indices)
  return tf.SparseTensor(
      indices=sp_tensor.indices,
      values=sp_tensor.values * batch_vec,
      dense_shape=sp_tensor.dense_shape) 

def rescore_sparse(sp_mentions, tf_db, qrys):
  """Rescore mentions in sparse tensor with given batch of queries."""
  batch_ind = sp_mentions.indices[:, 0]
  mention_ind = sp_mentions.indices[:, 1]
  mention_vec = tf.gather(tf_db, mention_ind)
  batch_qrys = tf.gather(qrys, batch_ind)
  mention_scs = tf.reduce_sum(batch_qrys * mention_vec, 1)
  return tf.SparseTensor(
      indices=sp_mentions.indices,
      values=mention_scs,
      dense_shape=sp_mentions.dense_shape) 

def dispatch(self, inp):
    """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>]`.
    """
    inp = tf.gather(inp, self._batch_index)
    return tf.split(inp, self._part_sizes_tensor, 0, num=self._num_experts) 

def gather_indexes(sequence_tensor, positions):
  """Gathers the vectors at the specific positions over a minibatch."""
  sequence_shape = modeling.get_shape_list(sequence_tensor, expected_rank=3)
  batch_size = sequence_shape[0]
  seq_length = sequence_shape[1]
  width = sequence_shape[2]

  flat_offsets = tf.reshape(
      tf.range(0, batch_size, dtype=tf.int32) * seq_length, [-1, 1])
  flat_positions = tf.reshape(positions + flat_offsets, [-1])
  flat_sequence_tensor = tf.reshape(sequence_tensor,
                                    [batch_size * seq_length, width])
  output_tensor = tf.gather(flat_sequence_tensor, flat_positions)
  return output_tensor 

def get_target_tokens_for_apply(token_ids, mask_indices):
  return tf.gather(token_ids, mask_indices) 

def mrcnn_class_loss_graph(target_class_ids, pred_class_logits,
                           active_class_ids):
    """Loss for the classifier head of Mask RCNN.

    target_class_ids: [batch, num_rois]. Integer class IDs. Uses zero
        padding to fill in the array.
    pred_class_logits: [batch, num_rois, num_classes]
    active_class_ids: [batch, num_classes]. Has a value of 1 for
        classes that are in the dataset of the image, and 0
        for classes that are not in the dataset.
    """
    # During model building, Keras calls this function with
    # target_class_ids of type float32. Unclear why. Cast it
    # to int to get around it.
    target_class_ids = tf.cast(target_class_ids, 'int64')

    # Find predictions of classes that are not in the dataset.
    pred_class_ids = tf.argmax(pred_class_logits, axis=2)
    # TODO: Update this line to work with batch > 1. Right now it assumes all
    #       images in a batch have the same active_class_ids
    pred_active = tf.gather(active_class_ids[0], pred_class_ids)

    # Loss
    loss = tf.nn.sparse_softmax_cross_entropy_with_logits(
        labels=target_class_ids, logits=pred_class_logits)

    # Erase losses of predictions of classes that are not in the active
    # classes of the image.
    loss = loss * pred_active

    # Computer loss mean. Use only predictions that contribute
    # to the loss to get a correct mean.
    loss = tf.reduce_sum(loss) / tf.reduce_sum(pred_active)
    return loss