jiahuimbzuai/CoQuIR huggingface dataset:
//  [python]
def testNotGreedyBeamTwoWithAlpha(self):
    batch_size = 1
    beam_size = 2
    vocab_size = 3
    decode_length = 3

    initial_ids = tf.constant([0] * batch_size)  # GO
    # Probabilities for position * batch * beam * vocab
    # Probabilities have been set such that with alpha = 3.5, the less probable
    # but longer sequence will have a better score that the shorter sequence
    # with higher log prob.
    probabilities = tf.constant([[[0.1, 0.1, 0.8], [0.1, 0.1, 0.8]],
                                 [[0.4, 0.5, 0.1], [0.2, 0.4, 0.4]],
                                 [[0.05, 0.9, 0.05], [0.4, 0.4, 0.2]]])

    def symbols_to_logits(ids):
      pos = tf.shape(ids)[1]
      logits = tf.cast(tf.log(probabilities[pos - 1, :]), tf.float32)
      return logits

    # Disable early stopping
    final_ids, final_scores, _ = beam_search.beam_search(
        symbols_to_logits,
        initial_ids,
        beam_size,
        decode_length,
        vocab_size,
        3.5,
        eos_id=1)

    with self.test_session():
      ids = final_ids.eval()
      scores = final_scores.eval()
    self.assertAllClose([[
        np.log(0.8 * 0.4 * 0.9) / (8. / 6.)**3.5,
        np.log(0.8 * 0.5) / (7. / 6.)**3.5
    ]], scores)
    self.assertAllEqual([[[0, 2, 0, 1], [0, 2, 1, 0]]], ids)
//  [python]
def testNotGreedyBatchTwoBeamTwoWithAlpha(self):
    batch_size = 2
    beam_size = 2
    vocab_size = 3
    decode_length = 3

    initial_ids = tf.constant([0] * batch_size)  # GO
    # Probabilities for position * batch * beam * vocab
    # Probabilities have been set such that with alpha = 3.5, the less probable
    # but longer sequence will have a better score than the shorter sequence
    # with higher log prob in batch 1, and the order will be reverse in batch
    # 2. That is, the shorter sequence will still have a higher score in spite
    # of the length penalty
    probabilities = tf.constant([[[[0.1, 0.1, 0.8], [0.1, 0.1, 0.8]],
                                  [[0.1, 0.1, 0.8], [0.1, 0.1, 0.8]]],
                                 [[[0.4, 0.5, 0.1], [0.2, 0.4, 0.4]],
                                  [[0.3, 0.6, 0.1], [0.2, 0.4, 0.4]]],
                                 [[[0.05, 0.9, 0.05], [0.4, 0.4, 0.2]],
                                  [[0.05, 0.9, 0.05], [0.4, 0.4, 0.2]]]])

    def symbols_to_logits(ids):
      pos = tf.shape(ids)[1]
      logits = tf.to_float(tf.log(probabilities[pos - 1, :]))
      return logits

    final_ids, final_scores, _ = beam_search.beam_search(
        symbols_to_logits,
        initial_ids,
        beam_size,
        decode_length,
        vocab_size,
        3.5,
        eos_id=1)

    with self.test_session():
      ids = final_ids.eval()
      scores = final_scores.eval()
    self.assertAllEqual([[[0, 2, 0, 1], [0, 2, 1, 0]], [[0, 2, 1, 0],
                                                        [0, 2, 0, 1]]], ids)
    self.assertAllClose([[
        np.log(0.8 * 0.4 * 0.9) / (8. / 6.)**3.5,
        np.log(0.8 * 0.5) / (7. / 6.)**3.5
    ], [
        np.log(0.8 * 0.6) / (7. / 6.)**3.5,
        np.log(0.8 * 0.3 * 0.9) / (8. / 6.)**3.5
    ]], scores)
//  [python]
def testNotGreedyBatchTwoBeamTwoWithAlpha(self):
    batch_size = 2
    beam_size = 2
    vocab_size = 3
    decode_length = 3

    initial_ids = tf.constant([0] * batch_size)  # GO
    # Probabilities for position * batch * beam * vocab
    # Probabilities have been set such that with alpha = 3.5, the less probable
    # but longer sequence will have a better score than the shorter sequence
    # with higher log prob in batch 1, and the order will be reverse in batch
    # 2. That is, the shorter sequence will still have a higher score in spite
    # of the length penalty
    probabilities = tf.constant([[[[0.1, 0.1, 0.8], [0.1, 0.1, 0.8]],
                                  [[0.1, 0.1, 0.8], [0.1, 0.1, 0.8]]],
                                 [[[0.4, 0.5, 0.1], [0.2, 0.4, 0.4]],
                                  [[0.3, 0.6, 0.1], [0.2, 0.4, 0.4]]],
                                 [[[0.05, 0.9, 0.05], [0.4, 0.4, 0.2]],
                                  [[0.05, 0.9, 0.05], [0.4, 0.4, 0.2]]]])

    def symbols_to_logits(ids):
      pos = tf.shape(ids)[1]
      logits = tf.cast(tf.log(probabilities[pos - 1, :]), tf.float32)
      return logits

    final_ids, final_scores, _ = beam_search.beam_search(
        symbols_to_logits,
        initial_ids,
        beam_size,
        decode_length,
        vocab_size,
        3.5,
        eos_id=1)

    with self.test_session():
      ids = final_ids.eval()
      scores = final_scores.eval()
    self.assertAllEqual([[[0, 2, 0, 1], [0, 2, 1, 0]], [[0, 2, 1, 0],
                                                        [0, 2, 0, 1]]], ids)
    self.assertAllClose([[
        np.log(0.8 * 0.4 * 0.9) / (8. / 6.)**3.5,
        np.log(0.8 * 0.5) / (7. / 6.)**3.5
    ], [
        np.log(0.8 * 0.6) / (7. / 6.)**3.5,
        np.log(0.8 * 0.3 * 0.9) / (8. / 6.)**3.5
    ]], scores)
//  [python]
def testStates(self):
    batch_size = 1
    beam_size = 1
    vocab_size = 2
    decode_length = 3

    initial_ids = tf.constant([0] * batch_size)  # GO
    probabilities = tf.constant([[[0.7, 0.3]], [[0.4, 0.6]], [[0.5, 0.5]]])

    expected_states = tf.constant([[[0.]], [[1.]]])

    def symbols_to_logits(ids, _, states):
      pos = tf.shape(ids)[1] - 1
      # We have to assert the values of state inline here since we can't fetch
      # them out of the loop!
      with tf.control_dependencies(
          [tf.assert_equal(states["state"], expected_states[pos])]):
        logits = tf.to_float(tf.log(probabilities[pos, :]))

      states["state"] += 1
      return logits, states

    states = {
        "state": tf.zeros((batch_size, 1)),
    }
    states["state"] = tf.placeholder_with_default(
        states["state"], shape=(None, 1))

    final_ids, _, _ = beam_search.beam_search(
        symbols_to_logits,
        initial_ids,
        beam_size,
        decode_length,
        vocab_size,
        0.0,
        eos_id=1,
        states=states)

    with self.test_session() as sess:
      # Catch and fail so that the testing framework doesn't think it's an error
      try:
        sess.run(final_ids)
      except tf.errors.InvalidArgumentError as e:
        raise AssertionError(e.message)
//  [python]
def testStates(self):
    batch_size = 1
    beam_size = 1
    vocab_size = 2
    decode_length = 3

    initial_ids = tf.constant([0] * batch_size)  # GO
    probabilities = tf.constant([[[0.7, 0.3]], [[0.4, 0.6]], [[0.5, 0.5]]])

    expected_states = tf.constant([[[0.]], [[1.]]])

    def symbols_to_logits(ids, _, states):
      pos = tf.shape(ids)[1] - 1
      # We have to assert the values of state inline here since we can't fetch
      # them out of the loop!
      with tf.control_dependencies(
          [tf.assert_equal(states["state"], expected_states[pos])]):
        logits = tf.cast(tf.log(probabilities[pos, :]), tf.float32)

      states["state"] += 1
      return logits, states

    states = {
        "state": tf.zeros((batch_size, 1)),
    }
    states["state"] = tf.placeholder_with_default(
        states["state"], shape=(None, 1))

    final_ids, _, _ = beam_search.beam_search(
        symbols_to_logits,
        initial_ids,
        beam_size,
        decode_length,
        vocab_size,
        0.0,
        eos_id=1,
        states=states)

    with self.test_session() as sess:
      # Catch and fail so that the testing framework doesn't think it's an error
      try:
        sess.run(final_ids)
      except tf.errors.InvalidArgumentError as e:
        raise AssertionError(e.message)
//  [python]
def testStatesAfterLoop(self):
    batch_size = 1
    beam_size = 1
    vocab_size = 2
    decode_length = 3

    initial_ids = tf.constant([0] * batch_size)  # GO
    probabilities = tf.constant([[[0.7, 0.3]], [[0.4, 0.6]], [[0.5, 0.5]]])

    def symbols_to_logits(ids, _, states):
      pos = tf.shape(ids)[1] - 1
      logits = tf.to_float(tf.log(probabilities[pos, :]))
      states["state"] += 1
      return logits, states

    states = {
        "state": tf.zeros((batch_size, 1)),
    }
    states["state"] = tf.placeholder_with_default(
        states["state"], shape=(None, 1))

    _, _, final_states = beam_search.beam_search(
        symbols_to_logits,
        initial_ids,
        beam_size,
        decode_length,
        vocab_size,
        0.0,
        eos_id=1,
        states=states)

    with self.test_session() as sess:
      final_states = sess.run(final_states)
    self.assertAllEqual([[[2]]], final_states["state"])
//  [python]
def testStatesAfterLoop(self):
    batch_size = 1
    beam_size = 1
    vocab_size = 2
    decode_length = 3

    initial_ids = tf.constant([0] * batch_size)  # GO
    probabilities = tf.constant([[[0.7, 0.3]], [[0.4, 0.6]], [[0.5, 0.5]]])

    def symbols_to_logits(ids, _, states):
      pos = tf.shape(ids)[1] - 1
      logits = tf.cast(tf.log(probabilities[pos, :]), tf.float32)
      states["state"] += 1
      return logits, states

    states = {
        "state": tf.zeros((batch_size, 1)),
    }
    states["state"] = tf.placeholder_with_default(
        states["state"], shape=(None, 1))

    _, _, final_states = beam_search.beam_search(
        symbols_to_logits,
        initial_ids,
        beam_size,
        decode_length,
        vocab_size,
        0.0,
        eos_id=1,
        states=states)

    with self.test_session() as sess:
      final_states = sess.run(final_states)
    self.assertAllEqual([[[2]]], final_states["state"])
//  [python]
def huber(x, kappa=0.1):
  return (0.5 * tf.square(x) * tf.to_float(tf.abs(x) <= kappa) +
          kappa * (tf.abs(x) - 0.5 * kappa) * tf.to_float(tf.abs(x) > kappa)
          ) / kappa
//  [python]
def huber(x, kappa=0.1):
  return (0.5 * tf.square(x) * tf.cast(tf.abs(x) <= kappa, tf.float32) +
          kappa * (tf.abs(x) - 0.5 * kappa) * tf.cast(tf.abs(x) > kappa, tf.float32)
          ) / kappa
//  [python]
def testStateBeamTwo(self):
    batch_size = 1
    beam_size = 2
    vocab_size = 3
    decode_length = 3

    initial_ids = tf.constant([0] * batch_size)  # GO
    probabilities = tf.constant([[[0.1, 0.1, 0.8], [0.1, 0.1, 0.8]],
                                 [[0.4, 0.5, 0.1], [0.2, 0.4, 0.4]],
                                 [[0.05, 0.9, 0.05], [0.4, 0.4, 0.2]]])

    # The top beam is always selected so we should see the top beam's state
    # at each position, which is the one thats getting 3 added to it each step.
    expected_states = tf.constant([[[0.], [0.]], [[3.], [3.]], [[6.], [6.]]])

    def symbols_to_logits(ids, _, states):
      pos = tf.shape(ids)[1] - 1

      # We have to assert the values of state inline here since we can't fetch
      # them out of the loop!
      with tf.control_dependencies(
          [tf.assert_equal(states["state"], expected_states[pos])]):
        logits = tf.to_float(tf.log(probabilities[pos, :]))

      states["state"] += tf.constant([[3.], [7.]])
      return logits, states

    states = {
        "state": tf.zeros((batch_size, 1)),
    }
    states["state"] = tf.placeholder_with_default(
        states["state"], shape=(None, 1))

    final_ids, _, _ = beam_search.beam_search(
        symbols_to_logits,
        initial_ids,
        beam_size,
        decode_length,
        vocab_size,
        0.0,
        eos_id=1,
        states=states)

    with self.test_session() as sess:
      # Catch and fail so that the testing framework doesn't think it's an error
      try:
        sess.run(final_ids)
      except tf.errors.InvalidArgumentError as e:
        raise AssertionError(e.message)
//  [python]
def testStateBeamTwo(self):
    batch_size = 1
    beam_size = 2
    vocab_size = 3
    decode_length = 3

    initial_ids = tf.constant([0] * batch_size)  # GO
    probabilities = tf.constant([[[0.1, 0.1, 0.8], [0.1, 0.1, 0.8]],
                                 [[0.4, 0.5, 0.1], [0.2, 0.4, 0.4]],
                                 [[0.05, 0.9, 0.05], [0.4, 0.4, 0.2]]])

    # The top beam is always selected so we should see the top beam's state
    # at each position, which is the one thats getting 3 added to it each step.
    expected_states = tf.constant([[[0.], [0.]], [[3.], [3.]], [[6.], [6.]]])

    def symbols_to_logits(ids, _, states):
      pos = tf.shape(ids)[1] - 1

      # We have to assert the values of state inline here since we can't fetch
      # them out of the loop!
      with tf.control_dependencies(
          [tf.assert_equal(states["state"], expected_states[pos])]):
        logits = tf.cast(tf.log(probabilities[pos, :]), tf.float32)

      states["state"] += tf.constant([[3.], [7.]])
      return logits, states

    states = {
        "state": tf.zeros((batch_size, 1)),
    }
    states["state"] = tf.placeholder_with_default(
        states["state"], shape=(None, 1))

    final_ids, _, _ = beam_search.beam_search(
        symbols_to_logits,
        initial_ids,
        beam_size,
        decode_length,
        vocab_size,
        0.0,
        eos_id=1,
        states=states)

    with self.test_session() as sess:
      # Catch and fail so that the testing framework doesn't think it's an error
      try:
        sess.run(final_ids)
      except tf.errors.InvalidArgumentError as e:
        raise AssertionError(e.message)
//  [python]
def testTPUBeam(self):
    batch_size = 1
    beam_size = 2
    vocab_size = 3
    decode_length = 3

    initial_ids = tf.constant([0] * batch_size)  # GO
    probabilities = tf.constant([[[0.1, 0.1, 0.8], [0.1, 0.1, 0.8]],
                                 [[0.4, 0.5, 0.1], [0.2, 0.4, 0.4]],
                                 [[0.05, 0.9, 0.05], [0.4, 0.4, 0.2]]])

    # The top beam is always selected so we should see the top beam's state
    # at each position, which is the one thats getting 3 added to it each step.
    expected_states = tf.constant([[[0.], [0.]], [[3.], [3.]], [[6.], [6.]]])

    def symbols_to_logits(_, i, states):
      # We have to assert the values of state inline here since we can't fetch
      # them out of the loop!
      with tf.control_dependencies(
          [tf.assert_equal(states["state"], expected_states[i])]):
        logits = tf.to_float(tf.log(probabilities[i, :]))

      states["state"] += tf.constant([[3.], [7.]])
      return logits, states

    states = {
        "state": tf.zeros((batch_size, 1)),
    }
    states["state"] = tf.placeholder_with_default(
        states["state"], shape=(None, 1))

    final_ids, _, _ = beam_search.beam_search(
        symbols_to_logits,
        initial_ids,
        beam_size,
        decode_length,
        vocab_size,
        3.5,
        eos_id=1,
        states=states,
        use_tpu=True)

    with self.test_session() as sess:
      # Catch and fail so that the testing framework doesn't think it's an error
      try:
        sess.run(final_ids)
      except tf.errors.InvalidArgumentError as e:
        raise AssertionError(e.message)
    self.assertAllEqual([[[0, 2, 0, 1], [0, 2, 1, 0]]], final_ids)
//  [python]
def testTPUBeam(self):
    batch_size = 1
    beam_size = 2
    vocab_size = 3
    decode_length = 3

    initial_ids = tf.constant([0] * batch_size)  # GO
    probabilities = tf.constant([[[0.1, 0.1, 0.8], [0.1, 0.1, 0.8]],
                                 [[0.4, 0.5, 0.1], [0.2, 0.4, 0.4]],
                                 [[0.05, 0.9, 0.05], [0.4, 0.4, 0.2]]])

    # The top beam is always selected so we should see the top beam's state
    # at each position, which is the one thats getting 3 added to it each step.
    expected_states = tf.constant([[[0.], [0.]], [[3.], [3.]], [[6.], [6.]]])

    def symbols_to_logits(_, i, states):
      # We have to assert the values of state inline here since we can't fetch
      # them out of the loop!
      with tf.control_dependencies(
          [tf.assert_equal(states["state"], expected_states[i])]):
        logits = tf.cast(tf.log(probabilities[i, :]), tf.float32)

      states["state"] += tf.constant([[3.], [7.]])
      return logits, states

    states = {
        "state": tf.zeros((batch_size, 1)),
    }
    states["state"] = tf.placeholder_with_default(
        states["state"], shape=(None, 1))

    final_ids, _, _ = beam_search.beam_search(
        symbols_to_logits,
        initial_ids,
        beam_size,
        decode_length,
        vocab_size,
        3.5,
        eos_id=1,
        states=states,
        use_tpu=True)

    with self.test_session() as sess:
      # Catch and fail so that the testing framework doesn't think it's an error
      try:
        sess.run(final_ids)
      except tf.errors.InvalidArgumentError as e:
        raise AssertionError(e.message)
    self.assertAllEqual([[[0, 2, 0, 1], [0, 2, 1, 0]]], final_ids)
//  [python]
def preprocess_example(self, example, mode, hparams):
    example = super(AudioTimitProblem, self).preprocess_example(
        example, mode, hparams)
    # Reshape audio to proper shape
    sample_count = tf.to_int32(example.pop("audio/sample_count"))
    sample_width = tf.to_int32(example.pop("audio/sample_width"))
    channel_count = 1
    example["inputs"] = tf.reshape(example["inputs"],
                                   [sample_count, sample_width, channel_count])
    return example
//  [python]
def preprocess_example(self, example, mode, hparams):
    example = super(AudioTimitProblem, self).preprocess_example(
        example, mode, hparams)
    # Reshape audio to proper shape
    sample_count = tf.cast(example.pop("audio/sample_count"), tf.int32)
    sample_width = tf.cast(example.pop("audio/sample_width"), tf.int32)
    channel_count = 1
    example["inputs"] = tf.reshape(example["inputs"],
                                   [sample_count, sample_width, channel_count])
    return example
//  [python]
def _format_groundtruth_data(self, image_shape):
    """Helper function for preparing groundtruth data for target assignment.

    In order to be consistent with the model.DetectionModel interface,
    groundtruth boxes are specified in normalized coordinates and classes are
    specified as label indices with no assumed background category.  To prepare
    for target assignment, we:
    1) convert boxes to absolute coordinates,
    2) add a background class at class index 0

    Args:
      image_shape: A 1-D int32 tensor of shape [4] representing the shape of the
        input image batch.

    Returns:
      groundtruth_boxlists: A list of BoxLists containing (absolute) coordinates
        of the groundtruth boxes.
      groundtruth_classes_with_background_list: A list of 2-D one-hot
        (or k-hot) tensors of shape [num_boxes, num_classes+1] containing the
        class targets with the 0th index assumed to map to the background class.
    """
    groundtruth_boxlists = [
        box_list_ops.to_absolute_coordinates(
            box_list.BoxList(boxes), image_shape[1], image_shape[2])
        for boxes in self.groundtruth_lists(fields.BoxListFields.boxes)]
    groundtruth_classes_with_background_list = [
        tf.to_float(
            tf.pad(one_hot_encoding, [[0, 0], [1, 0]], mode='CONSTANT'))
        for one_hot_encoding in self.groundtruth_lists(
            fields.BoxListFields.classes)]
    return groundtruth_boxlists, groundtruth_classes_with_background_list
//  [python]
def _format_groundtruth_data(self, image_shape):
    """Helper function for preparing groundtruth data for target assignment.

    In order to be consistent with the model.DetectionModel interface,
    groundtruth boxes are specified in normalized coordinates and classes are
    specified as label indices with no assumed background category.  To prepare
    for target assignment, we:
    1) convert boxes to absolute coordinates,
    2) add a background class at class index 0

    Args:
      image_shape: A 1-D int32 tensor of shape [4] representing the shape of the
        input image batch.

    Returns:
      groundtruth_boxlists: A list of BoxLists containing (absolute) coordinates
        of the groundtruth boxes.
      groundtruth_classes_with_background_list: A list of 2-D one-hot
        (or k-hot) tensors of shape [num_boxes, num_classes+1] containing the
        class targets with the 0th index assumed to map to the background class.
    """
    groundtruth_boxlists = [
        box_list_ops.to_absolute_coordinates(
            box_list.BoxList(boxes), image_shape[1], image_shape[2])
        for boxes in self.groundtruth_lists(fields.BoxListFields.boxes)]
    groundtruth_classes_with_background_list = [
        tf.cast(
            tf.pad(one_hot_encoding, [[0, 0], [1, 0]], mode='CONSTANT'), dtype=tf.float32)
        for one_hot_encoding in self.groundtruth_lists(
            fields.BoxListFields.classes)]
    return groundtruth_boxlists, groundtruth_classes_with_background_list
//  [python]
def state_preprocess_net(
    states,
    num_output_dims=2,
    states_hidden_layers=(100,),
    normalizer_fn=None,
    activation_fn=tf.nn.relu,
    zero_time=True,
    images=False):
  """Creates a simple feed forward net for embedding states.
  """
  with slim.arg_scope(
      [slim.fully_connected],
      activation_fn=activation_fn,
      normalizer_fn=normalizer_fn,
      weights_initializer=slim.variance_scaling_initializer(
          factor=1.0/3.0, mode='FAN_IN', uniform=True)):

    states_shape = tf.shape(states)
    states_dtype = states.dtype
    states = tf.to_float(states)
    if images:  # Zero-out x-y
      states *= tf.constant([0.] * 2 + [1.] * (states.shape[-1] - 2), dtype=states.dtype)
    if zero_time:
      states *= tf.constant([1.] * (states.shape[-1] - 1) + [0.], dtype=states.dtype)
    orig_states = states
    embed = states
    if states_hidden_layers:
      embed = slim.stack(embed, slim.fully_connected, states_hidden_layers,
                         scope='states')

    with slim.arg_scope([slim.fully_connected],
                        weights_regularizer=None,
                        weights_initializer=tf.random_uniform_initializer(
                            minval=-0.003, maxval=0.003)):
      embed = slim.fully_connected(embed, num_output_dims,
                                   activation_fn=None,
                                   normalizer_fn=None,
                                   scope='value')

    output = embed
    output = tf.cast(output, states_dtype)
    return output
//  [python]
def state_preprocess_net(
    states,
    num_output_dims=2,
    states_hidden_layers=(100,),
    normalizer_fn=None,
    activation_fn=tf.nn.relu,
    zero_time=True,
    images=False):
  """Creates a simple feed forward net for embedding states.
  """
  with slim.arg_scope(
      [slim.fully_connected],
      activation_fn=activation_fn,
      normalizer_fn=normalizer_fn,
      weights_initializer=slim.variance_scaling_initializer(
          factor=1.0/3.0, mode='FAN_IN', uniform=True)):

    states_shape = tf.shape(states)
    states_dtype = states.dtype
    states = tf.cast(states, tf.float32)
    if images:  # Zero-out x-y
      states *= tf.constant([0.] * 2 + [1.] * (states.shape[-1] - 2), dtype=states.dtype)
    if zero_time:
      states *= tf.constant([1.] * (states.shape[-1] - 1) + [0.], dtype=states.dtype)
    orig_states = states
    embed = states
    if states_hidden_layers:
      embed = slim.stack(embed, slim.fully_connected, states_hidden_layers,
                         scope='states')

    with slim.arg_scope([slim.fully_connected],
                        weights_regularizer=None,
                        weights_initializer=tf.random_uniform_initializer(
                            minval=-0.003, maxval=0.003)):
      embed = slim.fully_connected(embed, num_output_dims,
                                   activation_fn=None,
                                   normalizer_fn=None,
                                   scope='value')

    output = embed
    output = tf.cast(output, states_dtype)
    return output
//  [python]
def action_embed_net(
    actions,
    states=None,
    num_output_dims=2,
    hidden_layers=(400, 300),
    normalizer_fn=None,
    activation_fn=tf.nn.relu,
    zero_time=True,
    images=False):
  """Creates a simple feed forward net for embedding actions.
  """
  with slim.arg_scope(
      [slim.fully_connected],
      activation_fn=activation_fn,
      normalizer_fn=normalizer_fn,
      weights_initializer=slim.variance_scaling_initializer(
          factor=1.0/3.0, mode='FAN_IN', uniform=True)):

    actions = tf.to_float(actions)
    if states is not None:
      if images:  # Zero-out x-y
        states *= tf.constant([0.] * 2 + [1.] * (states.shape[-1] - 2), dtype=states.dtype)
      if zero_time:
        states *= tf.constant([1.] * (states.shape[-1] - 1) + [0.], dtype=states.dtype)
      actions = tf.concat([actions, tf.to_float(states)], -1)

    embed = actions
    if hidden_layers:
      embed = slim.stack(embed, slim.fully_connected, hidden_layers,
                         scope='hidden')

    with slim.arg_scope([slim.fully_connected],
                        weights_regularizer=None,
                        weights_initializer=tf.random_uniform_initializer(
                            minval=-0.003, maxval=0.003)):
      embed = slim.fully_connected(embed, num_output_dims,
                                   activation_fn=None,
                                   normalizer_fn=None,
                                   scope='value')
      if num_output_dims == 1:
        return embed[:, 0, ...]
      else:
        return embed
//  [python]
def action_embed_net(
    actions,
    states=None,
    num_output_dims=2,
    hidden_layers=(400, 300),
    normalizer_fn=None,
    activation_fn=tf.nn.relu,
    zero_time=True,
    images=False):
  """Creates a simple feed forward net for embedding actions.
  """
  with slim.arg_scope(
      [slim.fully_connected],
      activation_fn=activation_fn,
      normalizer_fn=normalizer_fn,
      weights_initializer=slim.variance_scaling_initializer(
          factor=1.0/3.0, mode='FAN_IN', uniform=True)):

    actions = tf.cast(actions, tf.float32)
    if states is not None:
      if images:  # Zero-out x-y
        states *= tf.constant([0.] * 2 + [1.] * (states.shape[-1] - 2), dtype=states.dtype)
      if zero_time:
        states *= tf.constant([1.] * (states.shape[-1] - 1) + [0.], dtype=states.dtype)
      actions = tf.concat([actions, tf.cast(states, tf.float32)], -1)

    embed = actions
    if hidden_layers:
      embed = slim.stack(embed, slim.fully_connected, hidden_layers,
                         scope='hidden')

    with slim.arg_scope([slim.fully_connected],
                        weights_regularizer=None,
                        weights_initializer=tf.random_uniform_initializer(
                            minval=-0.003, maxval=0.003)):
      embed = slim.fully_connected(embed, num_output_dims,
                                   activation_fn=None,
                                   normalizer_fn=None,
                                   scope='value')
      if num_output_dims == 1:
        return embed[:, 0, ...]
      else:
        return embed
//  [python]
def _example_index_to_sparse_index(example_indices, batch_size):
  """Creates a 2D sparse index tensor from a list of 1D example indices.

  For example, this would do the transformation:
  [0, 0, 0, 1, 3, 3] -> [[0,0], [0,1], [0,2], [1,0], [3,0], [3,1]]

  The second column of the output tensor is the running count of the occurrences
  of that example index.

  Args:
    example_indices: A sorted 1D Tensor with example indices.
    batch_size: The batch_size. Could be larger than max(example_indices) if the
      last examples of the batch do not have the feature present.

  Returns:
    The sparse index tensor.
    The maxmium length of a row in this tensor.
  """
  binned_counts = tf.bincount(example_indices, minlength=batch_size)
  max_len = tf.to_int64(tf.reduce_max(binned_counts))
  return tf.where(tf.sequence_mask(binned_counts)), max_len
//  [python]
def _example_index_to_sparse_index(example_indices, batch_size):
  """Creates a 2D sparse index tensor from a list of 1D example indices.

  For example, this would do the transformation:
  [0, 0, 0, 1, 3, 3] -> [[0,0], [0,1], [0,2], [1,0], [3,0], [3,1]]

  The second column of the output tensor is the running count of the occurrences
  of that example index.

  Args:
    example_indices: A sorted 1D Tensor with example indices.
    batch_size: The batch_size. Could be larger than max(example_indices) if the
      last examples of the batch do not have the feature present.

  Returns:
    The sparse index tensor.
    The maxmium length of a row in this tensor.
  """
  binned_counts = tf.bincount(example_indices, minlength=batch_size)
  max_len = tf.cast(tf.reduce_max(binned_counts), tf.int64)
  return tf.where(tf.sequence_mask(binned_counts)), max_len
//  [python]
def generate_landmarks(keypoints):
    is_background = tf.equal(keypoints, 0)
    ones = tf.to_float(tf.ones_like(is_background))
    zeros = tf.to_float(tf.zeros_like(is_background))

    return tf.where(is_background, zeros, ones) * 255
//  [python]
def generate_landmarks(keypoints):
    is_background = tf.equal(keypoints, 0)
    ones = tf.cast(tf.ones_like(is_background), tf.float32)
    zeros = tf.cast(tf.zeros_like(is_background), tf.float32)

    return tf.where(is_background, zeros, ones) * 255
//  [python]
def _dense_intersect_indices(tensor, required_sp_tensor):
  required_2d_indices = required_sp_tensor.indices[:, 0:2]
  values = tf.gather_nd(tensor, required_2d_indices)
  indices, max_len = _example_index_to_sparse_index(
      tf.to_int32(required_sp_tensor.indices[:, 0]),
      tf.to_int32(required_sp_tensor.dense_shape[0]))
  return tf.expand_dims(
      tf.sparse_to_dense(
          indices, tf.stack([required_sp_tensor.dense_shape[0], max_len]),
          values),
      axis=2)
//  [python]
def _dense_intersect_indices(tensor, required_sp_tensor):
  required_2d_indices = required_sp_tensor.indices[:, 0:2]
  values = tf.gather_nd(tensor, required_2d_indices)
  indices, max_len = _example_index_to_sparse_index(
      tf.cast(required_sp_tensor.indices[:, 0], tf.int32),
      tf.cast(required_sp_tensor.dense_shape[0], tf.int32))
  return tf.expand_dims(
      tf.sparse_to_dense(
          indices, tf.stack([required_sp_tensor.dense_shape[0], max_len]),
          values),
      axis=2)
//  [python]
def _remove_empty_timesteps(sp_tensor):
  """Creates a 3D SparseTensor skipping empty time steps.

  Args:
    sp_tensor: A SparseTensor with at least 2 dimensions (subsequent ones will
      be ignored and simply flattened into the 2nd dimension).

  Returns:
    A 3D SparseTensor with index 0 for dimension 3 and a series from [0,..k]
    for dimension 1 for each batch entry.
  """

  batch_size = tf.to_int32(sp_tensor.dense_shape[0])
  indices, max_len = _example_index_to_sparse_index(
      tf.to_int32(sp_tensor.indices[:, 0]), batch_size)
  indices = tf.concat([indices, tf.zeros_like(indices[:, 0:1])], axis=1)
  return tf.SparseTensor(
      indices=indices,
      values=sp_tensor.values,
      dense_shape=[batch_size, max_len, 1])
//  [python]
def _remove_empty_timesteps(sp_tensor):
  """Creates a 3D SparseTensor skipping empty time steps.

  Args:
    sp_tensor: A SparseTensor with at least 2 dimensions (subsequent ones will
      be ignored and simply flattened into the 2nd dimension).

  Returns:
    A 3D SparseTensor with index 0 for dimension 3 and a series from [0,..k]
    for dimension 1 for each batch entry.
  """

  batch_size = tf.cast(sp_tensor.dense_shape[0], tf.int32)
  indices, max_len = _example_index_to_sparse_index(
      tf.cast(sp_tensor.indices[:, 0], tf.int32), batch_size)
  indices = tf.concat([indices, tf.zeros_like(indices[:, 0:1])], axis=1)
  return tf.SparseTensor(
      indices=indices,
      values=sp_tensor.values,
      dense_shape=[batch_size, max_len, 1])
//  [python]
def keypts_encoding(keypoints, num_classes):
    keypoints = tf.to_int32(keypoints)
    keypoints = tf.reshape(keypoints, (-1,))
    keypoints = slim.layers.one_hot_encoding(keypoints, num_classes=num_classes+1)
    return keypoints
//  [python]
def keypts_encoding(keypoints, num_classes):
    keypoints = tf.cast(keypoints, tf.int32)
    keypoints = tf.reshape(keypoints, (-1,))
    keypoints = slim.layers.one_hot_encoding(keypoints, num_classes=num_classes+1)
    return keypoints
//  [python]
def _extend_with_dummy(extend_with, to_extend, dummy_value='n/a'):
  """Extends one SparseTensor with dummy_values at positions of other."""
  dense_shape = tf.to_int64(
      tf.concat([[tf.shape(extend_with)[0]],
                 [tf.maximum(tf.shape(extend_with)[1],
                             tf.shape(to_extend)[1])],
                 [tf.maximum(tf.shape(extend_with)[2],
                             tf.shape(to_extend)[2])]],
                axis=0))
  additional_indices = tf.sets.set_difference(
      tf.SparseTensor(
          indices=extend_with.indices,
          values=tf.zeros_like(extend_with.values, dtype=tf.int32),
          dense_shape=dense_shape),
      tf.SparseTensor(
          indices=to_extend.indices,
          values=tf.zeros([tf.shape(to_extend.indices)[0]], dtype=tf.int32),
          dense_shape=dense_shape)).indices
  # Supply defaults for all other indices.
  default = tf.tile(
      tf.constant([dummy_value]), multiples=[tf.shape(additional_indices)[0]])

  string_value = (
      tf.as_string(to_extend.values)
      if to_extend.values.dtype != tf.string else to_extend.values)
  return tf.sparse_reorder(
      tf.SparseTensor(
          indices=tf.concat([to_extend.indices, additional_indices], axis=0),
          values=tf.concat([string_value, default], axis=0),
          dense_shape=dense_shape))
//  [python]
def _extend_with_dummy(extend_with, to_extend, dummy_value='n/a'):
  """Extends one SparseTensor with dummy_values at positions of other."""
  dense_shape = tf.cast(
      tf.concat([[tf.shape(extend_with)[0]],
                 [tf.maximum(tf.shape(extend_with)[1],
                             tf.shape(to_extend)[1])],
                 [tf.maximum(tf.shape(extend_with)[2],
                             tf.shape(to_extend)[2])]],
                axis=0), tf.int64)
  additional_indices = tf.sets.set_difference(
      tf.SparseTensor(
          indices=extend_with.indices,
          values=tf.zeros_like(extend_with.values, dtype=tf.int32),
          dense_shape=dense_shape),
      tf.SparseTensor(
          indices=to_extend.indices,
          values=tf.zeros([tf.shape(to_extend.indices)[0]], dtype=tf.int32),
          dense_shape=dense_shape)).indices
  # Supply defaults for all other indices.
  default = tf.tile(
      tf.constant([dummy_value]), multiples=[tf.shape(additional_indices)[0]])

  string_value = (
      tf.as_string(to_extend.values)
      if to_extend.values.dtype != tf.string else to_extend.values)
  return tf.sparse_reorder(
      tf.SparseTensor(
          indices=tf.concat([to_extend.indices, additional_indices], axis=0),
          values=tf.concat([string_value, default], axis=0),
          dense_shape=dense_shape))
//  [python]
def get_weight(keypoints, mask=None, ng_w=0.01, ps_w=1.0):
    is_background = tf.equal(keypoints, 0)
    ones = tf.to_float(tf.ones_like(is_background))
    weights = tf.where(is_background, ones * ng_w, ones*ps_w)
    # if mask is not None:
    #     weights *= tf.to_float(mask)

    return weights
//  [python]
def get_weight(keypoints, mask=None, ng_w=0.01, ps_w=1.0):
    is_background = tf.equal(keypoints, 0)
    ones = tf.cast(tf.ones_like(is_background), tf.float32)
    weights = tf.where(is_background, ones * ng_w, ones*ps_w)
    # if mask is not None:
    #     weights *= tf.cast(mask, tf.float32)

    return weights
//  [python]
def forward(self, re_embedding, batch_size, train=True):
        rm = tf.reduce_sum(self.mask, 1)
        last_id = tf.gather_nd(self.alias, tf.stack([tf.range(batch_size), tf.to_int32(rm)-1], axis=1))
        last_h = tf.gather_nd(re_embedding, tf.stack([tf.range(batch_size), last_id], axis=1))
        seq_h = tf.stack([tf.nn.embedding_lookup(re_embedding[i], self.alias[i]) for i in range(batch_size)],
                         axis=0)                                                           #batch_size*T*d
        last = tf.matmul(last_h, self.nasr_w1)
        seq = tf.matmul(tf.reshape(seq_h, [-1, self.out_size]), self.nasr_w2)
        last = tf.reshape(last, [batch_size, 1, -1])
        m = tf.nn.sigmoid(last + tf.reshape(seq, [batch_size, -1, self.out_size]) + self.nasr_b)
        coef = tf.matmul(tf.reshape(m, [-1, self.out_size]), self.nasr_v, transpose_b=True) * tf.reshape(
            self.mask, [-1, 1])
        b = self.embedding[1:]
        if not self.nonhybrid:
            ma = tf.concat([tf.reduce_sum(tf.reshape(coef, [batch_size, -1, 1]) * seq_h, 1),
                            tf.reshape(last, [-1, self.out_size])], -1)
            self.B = tf.get_variable('B', [2 * self.out_size, self.out_size],
                                     initializer=tf.random_uniform_initializer(-self.stdv, self.stdv))
            y1 = tf.matmul(ma, self.B)
            logits = tf.matmul(y1, b, transpose_b=True)
        else:
            ma = tf.reduce_sum(tf.reshape(coef, [batch_size, -1, 1]) * seq_h, 1)
            logits = tf.matmul(ma, b, transpose_b=True)
        loss = tf.reduce_mean(
            tf.nn.sparse_softmax_cross_entropy_with_logits(labels=self.tar - 1, logits=logits))
        self.vars = tf.trainable_variables()
        if train:
            lossL2 = tf.add_n([tf.nn.l2_loss(v) for v in self.vars if v.name not
                               in ['bias', 'gamma', 'b', 'g', 'beta']]) * self.L2
            loss = loss + lossL2
        return loss, logits
//  [python]
def forward(self, re_embedding, batch_size, train=True):
        rm = tf.reduce_sum(self.mask, 1)
        last_id = tf.gather_nd(self.alias, tf.stack([tf.range(batch_size), tf.cast(rm-1, tf.int32)], axis=1))
        last_h = tf.gather_nd(re_embedding, tf.stack([tf.range(batch_size), last_id], axis=1))
        seq_h = tf.stack([tf.nn.embedding_lookup(re_embedding[i], self.alias[i]) for i in range(batch_size)],
                         axis=0)                                                           #batch_size*T*d
        last = tf.matmul(last_h, self.nasr_w1)
        seq = tf.matmul(tf.reshape(seq_h, [-1, self.out_size]), self.nasr_w2)
        last = tf.reshape(last, [batch_size, 1, -1])
        m = tf.nn.sigmoid(last + tf.reshape(seq, [batch_size, -1, self.out_size]) + self.nasr_b)
        coef = tf.matmul(tf.reshape(m, [-1, self.out_size]), self.nasr_v, transpose_b=True) * tf.reshape(
            self.mask, [-1, 1])
        b = self.embedding[1:]
        if not self.nonhybrid:
            ma = tf.concat([tf.reduce_sum(tf.reshape(coef, [batch_size, -1, 1]) * seq_h, 1),
                            tf.reshape(last, [-1, self.out_size])], -1)
            self.B = tf.get_variable('B', [2 * self.out_size, self.out_size],
                                     initializer=tf.random_uniform_initializer(-self.stdv, self.stdv))
            y1 = tf.matmul(ma, self.B)
            logits = tf.matmul(y1, b, transpose_b=True)
        else:
            ma = tf.reduce_sum(tf.reshape(coef, [batch_size, -1, 1]) * seq_h, 1)
            logits = tf.matmul(ma, b, transpose_b=True)
        loss = tf.reduce_mean(
            tf.nn.sparse_softmax_cross_entropy_with_logits(labels=self.tar - 1, logits=logits))
        self.vars = tf.trainable_variables()
        if train:
            lossL2 = tf.add_n([tf.nn.l2_loss(v) for v in self.vars if v.name not
                               in ['bias', 'gamma', 'b', 'g', 'beta']]) * self.L2
            loss = loss + lossL2
        return loss, logits
//  [python]
def file_based_input_fn_builder(input_file, batch_size, seq_length, is_training,
                                drop_remainder, hvd=None):
  """Creates an `input_fn` closure to be passed to Estimator."""

  name_to_features = {
      "input_ids": tf.io.FixedLenFeature([seq_length], tf.int64),
      "input_mask": tf.io.FixedLenFeature([seq_length], tf.int64),
      "segment_ids": tf.io.FixedLenFeature([seq_length], tf.int64),
      "label_ids": tf.io.FixedLenFeature([], tf.int64),
  }

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.to_int32(t)
      example[name] = t

    return example

  def input_fn():
    """The actual input function."""

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    d = tf.data.TFRecordDataset(input_file)
    if is_training:
      if hvd is not None: d = d.shard(hvd.size(), hvd.rank())
      d = d.repeat()
      d = d.shuffle(buffer_size=100)

    d = d.apply(
        tf.contrib.data.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder))

    return d

  return input_fn
//  [python]
def file_based_input_fn_builder(input_file, batch_size, seq_length, is_training,
                                drop_remainder, hvd=None):
  """Creates an `input_fn` closure to be passed to Estimator."""

  name_to_features = {
      "input_ids": tf.io.FixedLenFeature([seq_length], tf.int64),
      "input_mask": tf.io.FixedLenFeature([seq_length], tf.int64),
      "segment_ids": tf.io.FixedLenFeature([seq_length], tf.int64),
      "label_ids": tf.io.FixedLenFeature([], tf.int64),
  }

  def _decode_record(record, name_to_features):
    """Decodes a record to a TensorFlow example."""
    example = tf.parse_single_example(record, name_to_features)

    # tf.Example only supports tf.int64, but the TPU only supports tf.int32.
    # So cast all int64 to int32.
    for name in list(example.keys()):
      t = example[name]
      if t.dtype == tf.int64:
        t = tf.cast(t, tf.int32)
      example[name] = t

    return example

  def input_fn():
    """The actual input function."""

    # For training, we want a lot of parallel reading and shuffling.
    # For eval, we want no shuffling and parallel reading doesn't matter.
    d = tf.data.TFRecordDataset(input_file)
    if is_training:
      if hvd is not None: d = d.shard(hvd.size(), hvd.rank())
      d = d.repeat()
      d = d.shuffle(buffer_size=100)

    d = d.apply(
        tf.contrib.data.map_and_batch(
            lambda record: _decode_record(record, name_to_features),
            batch_size=batch_size,
            drop_remainder=drop_remainder))

    return d

  return input_fn
//  [python]
def _build_target_distribution(self):
    batch_size = tf.shape(self._replay.rewards)[0]
    # size of rewards: batch_size x 1
    rewards = self._replay.rewards[:, None]
    # size of tiled_support: batch_size x num_atoms

    is_terminal_multiplier = 1. - tf.cast(self._replay.terminals, tf.float32)
    # Incorporate terminal state to discount factor.
    # size of gamma_with_terminal: batch_size x 1
    gamma_with_terminal = self.cumulative_gamma * is_terminal_multiplier
    gamma_with_terminal = gamma_with_terminal[:, None]

    # size of next_qt_argmax: 1 x batch_size
    next_qt_argmax = tf.argmax(
        self._replay_next_target_net_outputs.q_values, axis=1)[:, None]
    batch_indices = tf.range(tf.to_int64(batch_size))[:, None]
    # size of next_qt_argmax: batch_size x 2
    batch_indexed_next_qt_argmax = tf.concat(
        [batch_indices, next_qt_argmax], axis=1)
    # size of next_logits (next quantiles): batch_size x num_atoms
    next_logits = tf.gather_nd(
        self._replay_next_target_net_outputs.logits,
        batch_indexed_next_qt_argmax)
    return rewards + gamma_with_terminal * next_logits
//  [python]
def _build_target_distribution(self):
    batch_size = tf.shape(self._replay.rewards)[0]
    # size of rewards: batch_size x 1
    rewards = self._replay.rewards[:, None]
    # size of tiled_support: batch_size x num_atoms

    is_terminal_multiplier = 1. - tf.cast(self._replay.terminals, tf.float32)
    # Incorporate terminal state to discount factor.
    # size of gamma_with_terminal: batch_size x 1
    gamma_with_terminal = self.cumulative_gamma * is_terminal_multiplier
    gamma_with_terminal = gamma_with_terminal[:, None]

    # size of next_qt_argmax: 1 x batch_size
    next_qt_argmax = tf.argmax(
        self._replay_next_target_net_outputs.q_values, axis=1)[:, None]
    batch_indices = tf.range(tf.cast(batch_size, tf.int64))[:, None]
    # size of next_qt_argmax: batch_size x 2
    batch_indexed_next_qt_argmax = tf.concat(
        [batch_indices, next_qt_argmax], axis=1)
    # size of next_logits (next quantiles): batch_size x num_atoms
    next_logits = tf.gather_nd(
        self._replay_next_target_net_outputs.logits,
        batch_indexed_next_qt_argmax)
    return rewards + gamma_with_terminal * next_logits
//  [python]
def attention_layer(embeds, nonzero_mask, opts):
    with tf.variable_scope('attention_layer', reuse=tf.AUTO_REUSE):
        attention_size = opts.attention_size
        # FC layer transform
        num_in = embeds.shape[2].value
        num_out = attention_size
        w = tf.sqrt(12.0 / (num_in + num_out))
        weights = tf.get_variable(
            'weights',
            initializer=tf.random_uniform([num_in, num_out], -w, w))
        biases = tf.get_variable(
            'biases',
            initializer=tf.zeros([num_out]))
        h = tf.tensordot(embeds, weights, 1) + biases
        # h = tf.nn.relu(h)

        # context vector
        u_context = tf.get_variable(
            'u_context',
            initializer=tf.random_uniform([attention_size], -w, w))

        hu = tf.reduce_sum(tf.multiply(h, u_context), axis=2)  # [batch, ws]
        tf.summary.histogram('hu', hu)
        hu_mask = tf.multiply(hu, nonzero_mask)  # [batch, ws]
        tf.summary.histogram('hu_mask', hu_mask)
        alphas = tf.nn.softmax(hu_mask)  # [batch, ws]
        tf.summary.histogram('alphas', alphas)

        tf.summary.histogram('nonzero_alphas',
                             tf.gather_nd(alphas, tf.where(alphas > 0.000001)))

        alpha_nonzeros_sum = tf.reduce_sum(tf.to_float(alphas > 0.0001))
        tf.summary.scalar('alpha_nonzeros_sum', alpha_nonzeros_sum)

        # output shape: [batch_size, embedding_size]
        alphas = tf.expand_dims(alphas, -1)  # [batch, ws, 1]
        attention_output = tf.reduce_sum(tf.multiply(embeds, alphas), axis=1)

    return attention_output
//  [python]
def attention_layer(embeds, nonzero_mask, opts):
    with tf.variable_scope('attention_layer', reuse=tf.AUTO_REUSE):
        attention_size = opts.attention_size
        # FC layer transform
        num_in = embeds.shape[2].value
        num_out = attention_size
        w = tf.sqrt(12.0 / (num_in + num_out))
        weights = tf.get_variable(
            'weights',
            initializer=tf.random_uniform([num_in, num_out], -w, w))
        biases = tf.get_variable(
            'biases',
            initializer=tf.zeros([num_out]))
        h = tf.tensordot(embeds, weights, 1) + biases
        # h = tf.nn.relu(h)

        # context vector
        u_context = tf.get_variable(
            'u_context',
            initializer=tf.random_uniform([attention_size], -w, w))

        hu = tf.reduce_sum(tf.multiply(h, u_context), axis=2)  # [batch, ws]
        tf.summary.histogram('hu', hu)
        hu_mask = tf.multiply(hu, nonzero_mask)  # [batch, ws]
        tf.summary.histogram('hu_mask', hu_mask)
        alphas = tf.nn.softmax(hu_mask)  # [batch, ws]
        tf.summary.histogram('alphas', alphas)

        tf.summary.histogram('nonzero_alphas',
                             tf.gather_nd(alphas, tf.where(alphas > 0.000001)))

        alpha_nonzeros_sum = tf.reduce_sum(tf.cast(alphas > 0.0001, tf.float32))
        tf.summary.scalar('alpha_nonzeros_sum', alpha_nonzeros_sum)

        # output shape: [batch_size, embedding_size]
        alphas = tf.expand_dims(alphas, -1)  # [batch, ws, 1]
        attention_output = tf.reduce_sum(tf.multiply(embeds, alphas), axis=1)

    return attention_output